Pytorch seems to freeze at "loss.item()"

Accelerated Computing CUDA CUDA Programming and Performance
1

Hello,

While training using pytorch with GPU, pytorch seems to freeze just after print " — 5."

Code fragment is as shown below.

loss = criterion(...)
optimizer.zero_grad()
loss.backward()
optimizer.step()

print('    --- 5', flush=True)
val = loss.item()
print('    --- 6', flush=True)

Is there anyone having same issue?
And, could anyone please tell me how to solve?

TIA