Hi,
These two lines look very suspicious:
        congestion_H.requires_grad = True
        congestion_V.requires_grad = True

Why do you have to do that? If the input features already require gradients, these outputs should as well.