update

Dockerfile

@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:1.12.0-cuda11.3-cudnn8-runtime
+FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime
 
 ARG USER
 ARG GROUP

utils.py

@@ -14,7 +14,8 @@ def imshow(tensor, figsize=None, title="", **args):
             imshow(el, figsize=figsize, title=title, **args)
             plt.suptitle("{} {}".format(idx, title))
         return
-    
+
+    print(type(tensor))
     tensor = tensor.detach().cpu() if type(tensor) == torch.Tensor else tensor
     if tensor.dtype == torch.complex64:
         f, ax = plt.subplots(1, 5, figsize=figsize, gridspec_kw={'width_ratios': [46.5,3,1,46.5,3]})