i’m trying to implemenet a training script with some modifications here’s what i’m doing:
i have image of (1, 2, 64, 64, 64) which i’m modifying to (1, 1, 64, 64, 64) with this:
It basically takes the largest probability at each point between the two channels in the output and then encodes it to either 0 or 1 in a new tensor.
def combine_channel(output):

  first, second = output[0][0], output[0][1]

  first = torch.reshape(first, shape=(1, 262144))
  second = torch.reshape(second, shape=(1, 262144))
  arr = []

  for fir, sec in zip(first[0], second[0]):
    if(fir > sec):
      val = 0
      arr.append(val)
    else:
      val = 1
      arr.append(val)

  

  final = torch.reshape(torch.from_numpy(np.asarray(arr)), shape=(1, 262144))
  image = torch.reshape(final, shape=(1, 1, 64, 64, 64)).type(torch.float32)
  return image

but then when I do this, it errors out with the error on the loss.backward() function:
epochs = 50
def train():
    criterion = DiceLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    num_steps_train = len(train_loader)
    
    print(num_steps_train)
    
    for epoch in range(1):
        t0 = time.time()
        for i, (images, masks) in enumerate(train_loader):
            images = images.to(device)
            masks = masks.to(device)
            outputs = model(images)
            outputs = combine_channel(outputs)
            outputs = outputs.to(device)
            
            print(outputs.size())
            print(masks.size())
            loss = criterion(outputs, masks)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            print('\r Epoch {} of {}   {:.2f}% - loss: {:.4f} - ({:.2f}s)'
                       .format(epoch + 1, epochs, i * 100 / num_steps_train, float
                               (loss), time.time() - t0))