Adds disable_autocast support for different device types.

modeling_phi.py

@@ -32,13 +32,16 @@ except:
     FusedDense = None
 
 
-def disable_autocast(func: Callable) -> Callable:
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        with torch.cuda.amp.autocast(enabled=False):
-            return func(*args, **kwargs)
-
-    return wrapper
+def disable_autocast(device_type: str = "cuda") -> None:
+    def _disable_autocast(f: Callable) -> Callable:
+        @wraps(f)
+        def __disable_autocast(*args, **kwargs) -> Callable:
+            with torch.autocast(device_type, enabled=False):
+                return f(*args, **kwargs)
+            
+        return __disable_autocast
+
+    return _disable_autocast
 
 
 @dataclass