I made the changes as you suggested. The grads are now being calculated, however all the gradients are zero, and the weights are becoming Nan. This happens immediately after the first backpropogation. This implicates that there is something wrong with the Loss function, but I have straight away implemented the loss function as it was mentioned in the paper. I doubt that there should be any problems with the non differentiability of the function.
class calculate_asn_loss_t_test(nn.Module):
    """This is a customized class to calculate the loss for the ASN block of the architecture
    using the t-test loss, please note that in all the tensors with the batch size B
    the first B/2 examples are the positive patches and the next B/2 samples are negative patches"""

    def __init__(self):
        super(calculate_asn_loss_t_test, self).__init__()


    def forward(self, original_tensor, reconstructed_tensor, beta=0.8, lambda_p=1, lambda_n=0.01):
        residual_l = original_tensor.size(0)
        residual_tensor = torch.abs(reconstructed_tensor - original_tensor).cuda()
        residual_tensor = residual_tensor.sum(dim=[1,2,3]).cuda()
        pos_residual = residual_tensor[:residual_l//2].float().cuda()
        neg_residual = residual_tensor[residual_l//2:].float().cuda()
        pos_mean = pos_residual.mean().cuda()
        neg_mean = neg_residual.mean().cuda()
        pos_std = pos_residual.std().cuda()
        neg_std = neg_residual.std().cuda()
        loss = torch.max(beta - pos_mean, torch.tensor([0.0]).cuda()) + neg_mean + lambda_p * (pos_std ** 2) + lambda_n * (neg_std ** 2)
        return loss