I downsized the network to the minimum required to reproduce the error. This code works fine on CPU with the same input shapes, but fails to work on GPU if two dimensions are >= 256 and the third >=4. 256x256x100 works on CPU without any issues. I should probably send a bug report on github
import torch
import torch.nn as nn
torch.set_default_tensor_type('torch.cuda.FloatTensor')

class Net(nn.Module):
    
    def __init__(self):
        super(Net,self).__init__()
        self.layer1=nn.Conv3d(1,64,5,padding=2)  #nn.BatchNorm3d(1)
        
        self.layer2 = nn.Conv3d(64,1,1) #binary mask classifier
        self.sigmoid=nn.Sigmoid()
            
            
    def forward(self,in_vol):
        
        first = self.layer1(in_vol)

        MonoClass=self.layer2(first)
        
        Mask=self.sigmoid(MonoClass)
        
        return Mask #torch.cat([Mask, Classes],dim=1) #return Mask, Classes


A=Net().cuda()
batch=1

SideSize=256
Zsize=4
X=torch.randn(batch,1,SideSize,SideSize,Zsize)
FakeMask=torch.randn(batch,1,SideSize,SideSize,Zsize).cuda()
FakeMask[FakeMask>0.5]=1
FakeMask[FakeMask<0.6]=0
X=X.cuda()

optimizer=torch.optim.Adam(A.parameters())

optimizer.zero_grad()
Mask = A(X)

loss=torch.sum(Mask)
loss=loss
loss.backward()
optimizer.step()
print(loss)