readme

src/main/inference.py

@@ -54,16 +54,23 @@ class Inference():
         self.all_frames = self.all_frames[:len(self.mel_chunk)]
         
         
-         # Image2Image translation model
-        self.image2image = ResUnetGenerator(input_nc=6,output_nc=3,num_downs=6,use_dropout=False)
+        # Image2Image translation model
+        self.image2image = ResUnetGenerator(input_nc=6,output_nc=3,num_downs=6,use_dropout=False).to(device)
 
         # Load pretrained weights to image2image model 
         image2image_weight = torch.load(self.image2image_ckpt, map_location=torch.device(device))['G']
+        
         # Since the checkpoint of model was trained using DataParallel with multiple GPU
         # It required to wrap a model with DataParallel wrapper class
-        self.image2image = DataParallel(self.image2image)
+        self.image2image = DataParallel(self.image2image).to(device)
         # assgin weight to model
         self.image2image.load_state_dict(image2image_weight)
+        
+        self.image2image = self.image2image.module # access model (remove DataParallel)
+        
+    
+        
+    
 
 
 
@@ -96,6 +103,8 @@ class Inference():
                         reset_optimizer=True,
                         pretrain=True)
 
+        print("Generator",next(self.generator.parameters()).is_cuda )
+        print("Img2Img",next(self.image2image.parameters()).is_cuda )
  
     def __landmark_detection__(self,images, batch_size):
         """
@@ -401,6 +410,7 @@ class Inference():
                     
                         with torch.no_grad():  
                             self.image2image.eval() 
+                            print("trans in", trans_in.is_cuda)
                             trans_out = self.image2image(trans_in)
                             trans_out = torch.tanh(trans_out)