Switch from torch.cuda.amp.custom_fwd to torch.amp.custom_fwd(device=...)

[loadams]

Fixes: #5682

[alexkirp]

xxx@As-MacBook-Pro ComfyUI % python3.12 main.py
Total VRAM 16384 MB, total RAM 16384 MB
pytorch version: 2.5.0.dev20240720
Set vram state to: SHARED
Device: mps
Using sub quadratic optimization for cross attention, if you have memory or speed issues try using: --use-split-cross-attention
[Prompt Server] web root: /Users/xxxstable_diffusion/ComfyUI/web
/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/kornia/feature/lightglue.py:44: FutureWarning: torch.cuda.amp.custom_fwd(args...) is deprecated. Please use torch.amp.custom_fwd(args..., device_type='cuda') instead.
@torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)

Import times for custom nodes:
0.0 seconds: /Users/akirpitchenko/stable_diffusion/ComfyUI/custom_nodes/websocket_image_save.py

Starting server

To see the GUI go to: http://127.0.0.1:8188

[loadams]

xxx@As-MacBook-Pro ComfyUI % python3.12 main.py Total VRAM 16384 MB, total RAM 16384 MB pytorch version: 2.5.0.dev20240720 Set vram state to: SHARED Device: mps Using sub quadratic optimization for cross attention, if you have memory or speed issues try using: --use-split-cross-attention [Prompt Server] web root: /Users/xxxstable_diffusion/ComfyUI/web /Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/kornia/feature/lightglue.py:44: FutureWarning: torch.cuda.amp.custom_fwd(args...) is deprecated. Please use torch.amp.custom_fwd(args..., device_type='cuda') instead. @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)

Import times for custom nodes: 0.0 seconds: /Users/akirpitchenko/stable_diffusion/ComfyUI/custom_nodes/websocket_image_save.py

Starting server

To see the GUI go to: http://127.0.0.1:8188

Hi @alexkirp - are you letting us know you hit the same issue on your mac? If so, thanks. I haven't had time to complete this PR yet, but we are working on it before it is fully deprecated.

[loadams]

Closing in favor of #5811