Hi Mukesh!
Three comments:



 Mukesh1729:

Input shape  = torch.rand(1,16,2,256,256) # 2 refers to the depth (in my case images)


First, if this really is the shape of the “image” you pass into your NET
model, it won’t work.  From your first post, the initial layer of NET is:
torch.nn.Conv3d(3, 64, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1))

That is, your Conv3d layer expects in_channels = 3, and you are
passing in a tensor with in_channels (the second dimension) equal
to 16.


def softXEnt (input, target):
    logprobs = torch.nn.functional.log_softmax (input, dim = 1)
    return  -(target * logprobs).sum() / input.shape[0]

input  = torch.rand(1, 16, 1, 256, 256)
target = torch.rand(1, 16, 1, 256, 256)
softXEnt(input, target)

output:tensor(1474121.)



Second, as written, the softXEnt() code I posted is taking the mean
over the batch (in your example, batch-size = 1), but is taking the sum
over the image dimensions (in your case, 1x256x256).
(The softXEnt() code I wrote wasn’t directed at the “multidimensional
case” with trailing “image” dimensions.)
If you prefer to average the loss over the image dimensions (which
might be more natural), you could do something like:
def softXEnt (input, target):
    logprobs = torch.nn.functional.log_softmax (input, dim = 1)
    return  -(target * logprobs).mean()

(I don’t know why I didn’t just use .mean() in the first place.)
Third, you are using input  = torch.rand(1, 16, 1, 256, 256).
I realize that this is just an example, but this input consists of values
that range over [0.0, 1.0] (because you are using rand()).  Such
values look like probabilities.  Make sure that your actual input to
softXEnt() (or CrossEntropyLoss, for that matter) consists of logits
that run from -inf to inf, and not probabilities.  (If you wanted to make
your example more realistic in this regard, you could use randn() in
place of rand().)
Best.
K. Frank