Solved by Ajinkya_Ambatwar in post #3 
              
              
                I created my custom implementation for Dijkstra using priority queues as discussed here 
For the same, I created a custom PriorityQ class using torch function as below 
class priorityQ_torch(object):
    """Priority Q implelmentation in PyTorch

    Args:
        object ([torch.Tensor]): [The Queue …