Zorobay:

grid = torch.stack((X.clone(), Y.clone(), torch.zeros_like(X)))
    grid = grid * torch.stack((a.clone(), b.clone(), torch.tensor(1.0))).view(3, 1, 1)



Give the above modifications a try. Let me know if this was the inplace operation