mirror of
https://github.com/vladmandic/automatic
synced 2026-09-13 01:59:42 +02:00
Cleanup
This commit is contained in:
+1
-1
@@ -665,6 +665,6 @@ def normalize_device(dev):
|
||||
|
||||
|
||||
def same_device(d1, d2):
|
||||
if d1.type != d2.type:
|
||||
if torch.device(d1).type != torch.device(d2).type:
|
||||
return False
|
||||
return normalize_device(d1) == normalize_device(d2)
|
||||
|
||||
@@ -36,7 +36,7 @@ def hijack_set_module_tensor(
|
||||
# note: majority of time is spent on .to(old_value.dtype)
|
||||
if tensor_name in module._buffers: # pylint: disable=protected-access
|
||||
module._buffers[tensor_name] = value.to(device, old_value.dtype) # pylint: disable=protected-access
|
||||
elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device): # pylint: disable=protected-access
|
||||
elif value is not None or not devices.same_device(device, module._parameters[tensor_name].device): # pylint: disable=protected-access
|
||||
param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access
|
||||
module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device, old_value.dtype) # pylint: disable=protected-access
|
||||
t1 = time.time()
|
||||
@@ -64,7 +64,7 @@ def hijack_set_module_tensor_simple(
|
||||
with devices.inference_context():
|
||||
if tensor_name in module._buffers: # pylint: disable=protected-access
|
||||
module._buffers[tensor_name] = value.to(device) # pylint: disable=protected-access
|
||||
elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device): # pylint: disable=protected-access
|
||||
elif value is not None or not devices.same_device(device, module._parameters[tensor_name].device): # pylint: disable=protected-access
|
||||
param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access
|
||||
module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device) # pylint: disable=protected-access
|
||||
t1 = time.time()
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user