change device type

handler.py

@@ -92,7 +92,7 @@ class EndpointHandler():
         
 
         # Define Generator with seed
-        self.generator = torch.Generator(device="cpu").manual_seed(3)
+        self.generator = torch.Generator(device=device.type).manual_seed(3)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """