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

@@ -15,6 +15,7 @@ def imshow(tensor, figsize=None, 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]})