add device specification

handler.py

@@ -27,7 +27,8 @@ class EndpointHandler:
         cfg = OmegaConf.create(config_dict)
 
         self.generator = Generator(cfg)
-        self.generator.load_state_dict(torch.load(generator_path))
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.generator.load_state_dict(torch.load(generator_path, map_location=self.device))
         self.generator.eval()
 
 
@@ -35,7 +36,7 @@ class EndpointHandler:
         base64_image = data.get('inputs')
         input_tensor = self._decode_base64_image(base64_image)
         print('Input tensor shape: ' + str(input_tensor.shape))
-        output_tensor = self.generator(input_tensor)
+        output_tensor = self.generator(input_tensor.to(self.device))
         output_tensor = output_tensor.squeeze(0)
         output_image = transforms.ToPILImage()(output_tensor)
         output_image = output_image.convert('RGB')