I am trying to perform some asynchronous processing where part of an algorithm is executed on an accelerator and then some processing is done on the GPU asynchronously.
I have a weird behavior involving F.mse_loss: I have commented out each line and it seems to cause a synchronization (defeating my goal of running async processing).
This is solved simply by switching the line
loss = F.mse_loss(y_hat, target)

with
loss = ((y_hat - target) ** 2).mean()

y_hat and target are CUDA Tensors of size (batch_size, 1)
I only call
loss.backward()
optimizer.step()
optimizer.zero_grad()

afterward and the sync issue goes away if I write the MSE explicitly.
I have tried to reproduce in an isolated way the problem, but running timings on F.mse_loss alone does not show this problem.
Is there any reason why a synchronization might be happening?