TypeError: Cannot handle this data type: (1, 1, 480, 640), <u2 about adabins HOT 5 OPEN

# predict depths of images stored in a directory and store the predictions in 16-bit format in a given separate dir
infer_helper.predict_dir("test_imgs", "test_imgs_results")  #test_imgs_results is a new directory to store output

KeyError                                  Traceback (most recent call last)
D:\anaconda3\lib\site-packages\PIL\Image.py in fromarray(obj, mode)
   2750         try:
-> 2751             mode, rawmode = _fromarray_typemap[typekey]
   2752         except KeyError as e:

KeyError: ((1, 1, 480, 640), '<u2')

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)
<ipython-input-7-36b5885687fe> in <module>
      1 # predict depths of images stored in a directory and store the predictions in 16-bit format in a given separate dir
----> 2 infer_helper.predict_dir("test_imgs", "test_imgs_results")

D:\anaconda3\lib\site-packages\torch\autograd\grad_mode.py in decorate_context(*args, **kwargs)
     24         def decorate_context(*args, **kwargs):
     25             with self.__class__():
---> 26                 return func(*args, **kwargs)
     27         return cast(F, decorate_context)
     28 

<private_repo>\infer.py in predict_dir(self, test_dir, out_dir)
    146             save_path = os.path.join(out_dir, basename + ".png")
    147 
--> 148             Image.fromarray(final.squeeze()).save(save_path)  #added .squeeze() to shrink size from 1,1,480,640 to 480,640
    149 
    150 

D:\anaconda3\lib\site-packages\PIL\Image.py in fromarray(obj, mode)
   2751             mode, rawmode = _fromarray_typemap[typekey]
   2752         except KeyError as e:
-> 2753             raise TypeError("Cannot handle this data type: %s, %s" % typekey) from e
   2754     else:
   2755         rawmode = mode

TypeError: Cannot handle this data type: (1, 1, 480, 640), <u2