Update app.py

app.py

@@ -144,9 +144,9 @@ async def split_model_weights():
         import torch
         import math
         
-        # Load the full model weights
+        # Load the full model weights without forcing CPU - let tensor servers handle device placement
         model_file = next(f for f in state.model_files.values() if f.endswith('.safetensors') or f.endswith('.bin'))
-        weights = torch.load(model_file, map_location='cpu')
+        weights = torch.load(model_file, weights_only=False)  # Explicitly allow non-weights for compatibility
         
         # Calculate total model size and chunks
         total_size_bytes = sum(p.nelement() * p.element_size() for p in weights.values())