Update tensor_server.py

tensor_server.py

@@ -191,16 +191,16 @@ def load_chunk(chunk: ModelChunk) -> torch.nn.Module:
                 return logits
         
                         
-                # Create and return the chunk buffer
-                chunk_model = ChunkBuffer(chunk_file, chunk_config)
-                # Ensure the chunk_model.config is the up-to-date config (including any assigned offsets)
-                chunk_model.config = chunk_config
-                print(f"[INFO] Loaded chunk {chunk.chunk_id} ({chunk_config.get('size_bytes', 0)} bytes) from {chunk.files[0]}")
-        
-                return chunk_model
+        # Create and return the chunk buffer
+        chunk_model = ChunkBuffer(chunk_file, chunk_config)
+        # Ensure the chunk_model.config is the up-to-date config (including any assigned offsets)
+        chunk_model.config = chunk_config
+        print(f"[INFO] Loaded chunk {chunk.chunk_id} ({chunk_config.get('size_bytes', 0)} bytes) from {chunk.files[0]}")
+
+        return chunk_model
                 
-            except Exception as e:
-                raise Exception(f"Failed to load chunk: {str(e)}")
+    except Exception as e:
+        raise Exception(f"Failed to load chunk: {str(e)}")
 
 async def process_tensor(chunk_id: int, inputs: torch.Tensor) -> torch.Tensor:
     """Process input tensor through the specified chunk"""