Mixed precision error

[guillaume-byte]

Seems like mixed precision is not working with the wrapped optimizer and model.

scaler.step(optimizer.get())
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/ubuntu/guillaume_plg/.venv/lib/python3.10/site-packages/torch/amp/grad_scaler.py", line 462, in step
    len(optimizer_state["found_inf_per_device"]) > 0
AssertionError: No inf checks were recorded for this optimizer.
scaler.step(optimizer.get().optimizer)
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/ubuntu/guillaume_plg/.venv/lib/python3.10/site-packages/torch/amp/grad_scaler.py", line 462, in step
    len(optimizer_state["found_inf_per_device"]) > 0
AssertionError: No inf checks were recorded for this optimizer.

[guillaume-byte]

Seems to be related only to torch / cuda local version