Palash_Shah:

outputs = combine_channel(outputs)


The issue is that you are converting tensors to numpy and it breaks the computational graph. I have provided a solution in your first code as there is easier way to achieve combine_channel.

  
    
    
    Torch tensor reshape by removing a channel vision
  
  
    You can do this by torch.max 
output = torch.randn(1, 2, 64, 64, 64)
output,_ = torch.max(output, 1, keepdim=True)
print(output.shape) # [1, 1, 64, 64, 64]