Bug&Fix due to Tensors Located on Different Files

handler.py

@@ -26,8 +26,7 @@ class EndpointHandler():
     def get_prediction(self, instance):
         instance["attention_mask"] = [[1] * len(instance["input_ids"])]
         for key in ["input_ids", "token_type_ids", "attention_mask"]:
-            instance[key] = torch.tensor(instance[key]).unsqueeze(0)  # Batch size = 1
-            instance[key].to(self.device)
+            instance[key] = torch.tensor(instance[key]).unsqueeze(0).to(self.device)  # Batch size = 1
 
         output = self.model(input_ids=instance["input_ids"],
                        attention_mask=instance["attention_mask"],