Upload models/pytorch_wrapper.py with huggingface_hub

models/pytorch_wrapper.py

@@ -66,6 +66,14 @@ class PyTorchModel:
             raise RuntimeError("Model not initialized.")
             
         try:
+            # Handle dictionary input (ONNX style)
+            if isinstance(input_data, dict):
+                # Take the first value (assuming single input) or look for "input"
+                if "input" in input_data:
+                    input_data = input_data["input"]
+                else:
+                    input_data = list(input_data.values())[0]
+
             # Convert to Tensor
             tensor_in = torch.from_numpy(input_data).float().to(self.device)