Hi,
Keep in mind that the GPU api is asynchronous. So you might want to add a torch.cuda.syncrhonize() after the line to make sure it finished executing.
Also since, you reuse the out variable, the old Tensor that out was pointing to is not reachable anymore and can be deleted when you’re not training (when training, it needs to be kept around to be able to compute the backward).