add option to return intermediates as tensors

climategan/trainer.py

@@ -334,19 +334,24 @@ class Trainer:
                     output_data["smog"] = smog
 
         if return_intermediates:
-            output_data["mask"] = (
-                ((mask > bin_value) * 255).cpu().numpy().astype(np.uint8)
-            )
-            output_data["depth"] = (
-                normalize_tensor(depth).cpu().squeeze(1).numpy().astype(np.uint8) * 255
-            )
-            output_data["segmentation"] = (
-                decode_segmap_merged_labels(segmentation, "r", False)
-                .cpu()
-                .permute(0, 2, 3, 1)
-                .numpy()
-                .astype(np.uint8)
-            )
+            if numpy:
+                output_data["mask"] = (
+                    ((mask > bin_value) * 255).cpu().numpy().astype(np.uint8)
+                )
+                output_data["depth"] = (
+                    normalize_tensor(depth).cpu().squeeze(1).numpy().astype(np.uint8) * 255
+                )
+                output_data["segmentation"] = (
+                    decode_segmap_merged_labels(segmentation, "r", False)
+                    .cpu()
+                    .permute(0, 2, 3, 1)
+                    .numpy()
+                    .astype(np.uint8)
+                )
+            else:
+                output_data["mask"] = mask
+                output_data["depth"] = depth
+                output_data["segmentation"] = segmentation
 
         return output_data