so i tried this by implementing my own function and it doesn’t seem to work:
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.float)
  return image

this just compares the two channels (since they’re probabilities of class 1 or 2) and then creates a new tensor with the larger one.