Hey Frank,
Thanks for getting back. Regarding the dimensions of the image that was a typo, but I have tried what you suggested and it works:
def softXEnt (input, target):
    logprobs = torch.nn.functional.log_softmax (input, dim = 1)
    return  -(target * logprobs).mean()

# Input image to the model: (1, 3, 2 ,256, 256) 

#Ouput from the model                  
model_out  = torch.randn(1, 16, 1 ,256, 256)
#Target heatmap
target_hmp = torch.randn(1, 16, 1 ,256, 256)

softXEnt(model_out, target_hmp)

output:tensor(-0.0006)

I added the above loss to my training loop and the loss looks like this. It starts off very low
I am using around 400 training images and 100 validation images
Legends:
blue: training
yellow:validation

Training loop
def loss_plot(epochs, train, val):
    ep = [i for i in range(epochs)]
    plt.plot(ep,train,label="Training loss")
    plt.plot(ep,val,label="Validation loss")
    plt.title("Training and validation loss")
    plt.show()
    
def softXEnt (input, target):
    logprobs = torch.nn.functional.log_softmax (input, dim = 1)
    return  -(target * logprobs).mean()



#training and Evaluation
#Get parameters, start training model
DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"
EPOCHS     = 10
lr         = 1e-3
model      = T_LEAP()
criterion  = torch.nn.MSELoss()
optimizer  = torch.optim.Adam(model.parameters(),lr=lr, amsgrad=True)

soft       = torch.nn.Softmax(dim=1)
# scheduler = StepLR(optimizer, step_size=5, gamma=0.1)


model = model.to(DEVICE)
tr_losses  = [] 
val_losses = []

tr_batch = int(len(train_loader)/train_loader.batch_size)
vl_batch = int(len(valid_loader)/valid_loader.batch_size)

for idx in range(EPOCHS):
  curr_trloss  = 0.0
  curr_valoss  = 0.0
  model.train()
    
  for i,d in enumerate(train_loader):
      
    # Get images, move them to GPU
    image   = d["images"]
    image   = image.permute(0,2,1,3,4).to(DEVICE)
    heatmap = d["heatmaps"]
    heatmap = heatmap.permute(0,4,1,2,3).to(DEVICE)
    

    optimizer.zero_grad()

    output = model(image)
    loss   = softXEnt(output, heatmap)
  
    loss.backward()
    optimizer.step()
    
    curr_trloss += loss.item()
  tr_losses.append(curr_trloss/tr_batch)
    
 
  # Evaluation loop
  model.eval()
  with torch.no_grad():
      for i,d in enumerate(valid_loader):

        # Get images, move them to GPU
        image   = d["images"]
        image   = image.permute(0,2,1,3,4).to(DEVICE)
        heatmap = d["heatmaps"]
        heatmap = heatmap.permute(0,4,1,2,3).to(DEVICE)

        output = model(image)
        loss   = softXEnt(output, heatmap)

        curr_valoss += loss.item()
        # scheduler.step()
          
  val_losses.append(curr_trloss/vl_batch)
    
  # Print losses 
  print(f"Training loss {idx}:",curr_trloss/tr_batch)
  print(f"Validation loss {idx}:",curr_valoss/vl_batch)
  print(" ")
    
print("Training and evalutation is now complete")
#plotloss    
loss_plot(EPOCHS,tr_losses, val_losses)

#save model
state = {
          'epoch': EPOCHS+1,
          'state_dict': model.state_dict(),
          'optimizer': optimizer.state_dict(),
          'learning_rate':lr,
        }
torch.save(state, "cow_model.pth")

My predictions versus ground truth look like:
Prediction

image1193×455 77.8 KB

Ground-truth

image1195×459 24.1 KB