Cleanup

2026-09-20 01:31:13 +02:00 · 2025-06-25 23:37:41 +03:00
parent e43d1d2ba7
commit c9f49720c5
3 changed files with 4 additions and 4 deletions
@@ -171,7 +171,7 @@ class OffloadHook(accelerate.hooks.ModelHook):
        return module

    def pre_forward(self, module, *args, **kwargs):
-        if devices.normalize_device(module.device) != devices.normalize_device(devices.device):
+        if not devices.same_device(module.device, devices.device):
            device_index = torch.device(devices.device).index
            if device_index is None:
                device_index = 0