fix for multi gpu, thanks to @TMullerSG

Author: lucidrains

med_seg_diff_pytorch/med_seg_diff_pytorch.py

@@ -494,12 +494,12 @@ def __init__(
     ):
         super().__init__()
 
-        self.model = model
+        self.model = model if isinstance(model, Unet) else model.module
+
         self.input_img_channels = self.model.input_img_channels
         self.mask_channels = self.model.mask_channels
         self.self_condition = self.model.self_condition
-
-        self.image_size = model.image_size
+        self.image_size = self.model.image_size
 
         self.objective = objective