Update app.py

app.py

@@ -144,9 +144,18 @@ async def split_model_weights():
         import torch
         import math
         
-        # Load the full model weights without forcing CPU - let tensor servers handle device placement
+        # Load the full model weights
         model_file = next(f for f in state.model_files.values() if f.endswith('.safetensors') or f.endswith('.bin'))
-        weights = torch.load(model_file, weights_only=False)  # Explicitly allow non-weights for compatibility
+        loaded_weights = torch.load(model_file, weights_only=False)  # Explicitly allow non-weights for compatibility
+        
+        # For OPT models, weights are stored under model.decoder
+        if isinstance(loaded_weights, dict):
+            if 'model.decoder' in str(list(loaded_weights.keys())):
+                # Get all weights that are part of the decoder
+                weights = {k: v for k, v in loaded_weights.items() if k.startswith('model.decoder')}
+            else:
+                # Just use all weights if no decoder prefix found
+                weights = loaded_weights
         
         # Calculate total model size and chunks
         total_size_bytes = sum(p.nelement() * p.element_size() for p in weights.values())