Competitive Collaborationで出力結果が保存できない

前提・実現したいこと

https://github.com/anuragranj/cc
のコードを実行中にエラーがでました。
python test_flow.py
--pretrained-disp /content/drive/'My Drive'/Competitive_Collaboration/pretrain/geometry/dispnet_cs_k.pth.tar
--pretrained-pose /content/drive/'My Drive'/Competitive_Collaboration/pretrain/geometry/posenet.pth.tar
--pretrained-mask /content/drive/'My Drive'/Competitive_Collaboration/pretrain/geometry/masknet.pth.tar
--pretrained-flow /content/drive/'My Drive'/Competitive_Collaboration/pretrain/geometry/back2future.pth.tar
--kitti-dir /content/drive/'My Drive'/Competitive_Collaboration/kitti2015
--output-dir /content/drive/'My Drive'/Competitive_Collaboration/cc/output
を実行しました。
最後のオプション --output-dir 以下をつけずに実行したときは正常動作していました。
ご多用のところ、恐れいりますが、ご教授いただけますと幸いです。

発生している問題・エラーメッセージ

/content/drive/My Drive/Competitive_Collaboration/cc
  0% 0/200 [00:00<?, ?it/s]test_flow.py:109: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.
  tgt_img_var = Variable(tgt_img.cuda(), volatile=True)
test_flow.py:110: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.
  ref_imgs_var = [Variable(img.cuda(), volatile=True) for img in ref_imgs]
test_flow.py:111: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.
  intrinsics_var = Variable(intrinsics.cuda(), volatile=True)
test_flow.py:112: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.
  intrinsics_inv_var = Variable(intrinsics_inv.cuda(), volatile=True)
test_flow.py:114: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.
  flow_gt_var = Variable(flow_gt.cuda(), volatile=True)
test_flow.py:115: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.
  obj_map_gt_var = Variable(obj_map_gt.cuda(), volatile=True)
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:2796: UserWarning: nn.functional.upsample is deprecated. Use nn.functional.interpolate instead.
  warnings.warn("nn.functional.upsample is deprecated. Use nn.functional.interpolate instead.")
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:2973: UserWarning: Default upsampling behavior when mode=bilinear is changed to align_corners=False since 0.4.0. Please specify align_corners=True if the old behavior is desired. See the documentation of nn.Upsample for details.
  "See the documentation of nn.Upsample for details.".format(mode))
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1569: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.
  warnings.warn("nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.")
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:3226: UserWarning: Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.
  warnings.warn("Default grid_sample and affine_grid behavior has changed "
  0% 0/200 [00:03<?, ?it/s]
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/PIL/Image.py", line 2680, in fromarray
    mode, rawmode = _fromarray_typemap[typekey]
KeyError: ((1, 1, 832), '|u1')

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "test_flow.py", line 211, in <module>
    main()
  File "test_flow.py", line 180, in main
    row1_viz_im = Image.fromarray((255*row1_viz).astype('uint8'))
  File "/usr/local/lib/python3.6/dist-packages/PIL/Image.py", line 2682, in fromarray
    raise TypeError("Cannot handle this data type: %s, %s" % typekey)
TypeError: Cannot handle this data type: (1, 1, 832), |u1

該当のソースコード

python
1            rigid_flow_viz = flow_to_image(tensor2array(flow_cam.data[0].cpu()))
2            non_rigid_flow_viz = flow_to_image(tensor2array(flow_fwd_non_rigid.data[0].cpu()))
3            total_flow_viz = flow_to_image(tensor2array(total_flow.data[0].cpu()))
4            row1_viz = np.hstack((tgt_img_viz, depth_viz, mask_viz))
5            row2_viz = np.hstack((rigid_flow_viz, non_rigid_flow_viz, total_flow_viz))
6
7            row1_viz_im = Image.fromarray((255*row1_viz).astype('uint8'))
8            row2_viz_im = Image.fromarray((row2_viz).astype('uint8'))
9
10            row1_viz_im.save(viz_dir/str(i).zfill(3)+'01.png')
11            row2_viz_im.save(viz_dir/str(i).zfill(3)+'02.png')