Fix Deprecation Warnings torch.GradScaler and torch.autocast

[IliasChair]

This PR replaces all occurrences of torch.cuda.amp.GradScaler(args...) and torch.cuda.amp.autocast(args...) with torch.GradScaler("cuda", args...) and torch.autocast("cuda", args...) respectively.

This fixes the deprecation warnings during training and inference like:

FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast("cuda", args...)` instead.
  with torch.cuda.amp.autocast(enabled=self.scaler is not None):

see PyTorch Docs

P.S. I'm not sure about the difference between torch.autocast("cuda", args...) and torch.amp.autocast("cuda", args...) (as seen in the deprecation warning). But I would to prefer to stick to the Documentation. E.g. only using torch.autocast without the amp

[codecov]

Codecov Report

Attention: Patch coverage is 80.00000% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...core/models/equiformer_v2/trainers/dens_trainer.py	0.00%	2 Missing ⚠️
...e/models/equiformer_v2/equiformer_v2_deprecated.py	0.00%	1 Missing ⚠️

Files with missing lines	Coverage Δ
...c/fairchem/core/models/equiformer_v2/layer_norm.py	57.89% <100.00%> (ø)
src/fairchem/core/models/gemnet_oc/gemnet_oc.py	89.91% <100.00%> (ø)
...rchem/core/models/gemnet_oc/layers/force_scaler.py	54.34% <ø> (ø)
src/fairchem/core/modules/scaling/fit.py	73.80% <100.00%> (ø)
src/fairchem/core/trainers/base_trainer.py	87.06% <100.00%> (ø)
src/fairchem/core/trainers/ocp_trainer.py	69.66% <100.00%> (ø)
...e/models/equiformer_v2/equiformer_v2_deprecated.py	91.15% <0.00%> (ø)
...core/models/equiformer_v2/trainers/dens_trainer.py	11.53% <0.00%> (ø)

[rayg1234]

@IliasChair thanks for helping us clean this up! This should be fine i would prefer to get the device instead blanket casting to cuda (there are tests that rely on cpu autocast) ie:
with torch.autocast(x.device, enabled=False)

but if the tests pass and the all the old code is using cuda.autocast anyways, it should be fine

[IliasChair]

Hi,
All the changes I made are in portions of the code that currently support GPU-based autocasting only. I don't see any existing code paths supporting CPU autocasting, so these changes should not introduce any issues.

If adding support for CPU autocasting is a planned feature, I believe it would still make sense to merge these changes first and then create a separate issue for that later.

As I see it, the codebase is largely tailored for GPU usage, particularly for more intensive tasks. Adding full CPU support would likely be a larger endavour, and it may not be worth the effort are going to use GPUs for anything machine learning related anyway.

That said, if I missed anything, I'd be happy to update my PR accordingly.

Best regards,
Ilias