Exception occured during image io error

Hi,

I have been trying to load the predicted segmentation mask, which is a 3d PyTorch tensor saved as a nifty file, but I am getting below error. I am actually able to load the image volumes stored in a similar manner but not sure what’s causing this issue while saving the segmentation mask.
image
Below is how I save the numpy array to nii.gz file.

       
        pred = nibabel.Nifti1Image(pred.numpy(), affine=np.eye(4))
       
        nibabel.save(pred, fnames[1])

What is dtype of pred.numpy()? Maybe pass pred.numpy().astype('uint8') instead?