diff options
| author | Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> | 2024-01-29 14:27:53 +0000 |
|---|---|---|
| committer | Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> | 2024-01-29 14:27:53 +0000 |
| commit | 750dd6014a45397979cad42a74634451d0861581 (patch) | |
| tree | a029d5a5f9321de1032d90a855ce38821282bca8 /modules/devices.py | |
| parent | 6484053037aaf13cd9774f150e8f45e11c8dc536 (diff) | |
| download | stable-diffusion-webui-gfx803-750dd6014a45397979cad42a74634451d0861581.tar.gz stable-diffusion-webui-gfx803-750dd6014a45397979cad42a74634451d0861581.tar.bz2 stable-diffusion-webui-gfx803-750dd6014a45397979cad42a74634451d0861581.zip | |
Fix potential bugs
Diffstat (limited to 'modules/devices.py')
| -rw-r--r-- | modules/devices.py | 9 |
1 files changed, 7 insertions, 2 deletions
diff --git a/modules/devices.py b/modules/devices.py index dfffaf24..60f7d6d7 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -141,7 +141,12 @@ def manual_cast_forward(target_dtype): args = [arg.to(target_dtype) if isinstance(arg, torch.Tensor) else arg for arg in args] kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()} - org_dtype = torch_utils.get_param(self).dtype + org_dtype = target_dtype + for param in self.parameters(): + if param.dtype != target_dtype: + org_dtype = param.dtype + break + if org_dtype != target_dtype: self.to(target_dtype) result = self.org_forward(*args, **kwargs) @@ -170,7 +175,7 @@ def manual_cast(target_dtype): continue applied = True org_forward = module_type.forward - if module_type == torch.nn.MultiheadAttention and has_xpu(): + if module_type == torch.nn.MultiheadAttention: module_type.forward = manual_cast_forward(torch.float32) else: module_type.forward = manual_cast_forward(target_dtype) |