Update inference.py

inference.py

@@ -6,8 +6,9 @@ def load_model():
     return model
 
 def predict(model, inputs):
-    inputs = [torch.tensor(input) for input in inputs]  # Batch inputs as tensors
+    inputs = [torch.tensor(input) for input in inputs]
     with torch.no_grad():
         output = model(*inputs)
-    return output.numpy().tolist()
+    return output.detach().cpu().numpy().tolist()
+