14 lines
287 B
Python
14 lines
287 B
Python
from contextlib import nullcontext
|
|
|
|
import torch
|
|
|
|
|
|
def autocast_exclude_mps(
|
|
device_type: str, dtype: torch.dtype
|
|
) -> nullcontext | torch.autocast:
|
|
return (
|
|
nullcontext()
|
|
if torch.backends.mps.is_available()
|
|
else torch.autocast(device_type, dtype)
|
|
)
|