Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve warning message when catching PyTorch MPS convolution bug #522

Merged
merged 1 commit into from
Dec 19, 2024
Merged

Improve warning message when catching PyTorch MPS convolution bug #522

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 22 additions & 10 deletions nam/models/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,10 @@ def _export_input_output(self) -> _Tuple[_np.ndarray, _np.ndarray]:
)


def _get_torch_version() -> str:
return _torch.__version__


class BaseNet(_Base):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
Expand Down Expand Up @@ -217,17 +221,25 @@ def _forward_mps_safe(self, x: _torch.Tensor, **kwargs) -> _torch.Tensor:
return self._forward(x, **kwargs)
except NotImplementedError as e:
if "Output channels > 65536 not supported at the MPS device." in str(e):
print(
"===WARNING===\n"
"NAM encountered a bug in PyTorch's MPS backend and will "
"switch to a fallback.\n"
f"Your version of PyTorch is {_torch.__version__}.\n"
"Please report this in an Issue at:\n"
"https://github.com/sdatkinson/neural-amp-modeler/issues/new/choose"
"\n"
"so that NAM's dependencies can avoid buggy versions of "
"PyTorch and the associated performance hit."
msg = (
"Warning: NAM encountered a bug in PyTorch's MPS backend and "
"will switch to a fallback."
)
known_bad_versions = {"2.5.0", "2.5.1"}
torch_version = _get_torch_version()
if torch_version not in known_bad_versions:
msg += (
"\n"
f"Your version of PyTorch is {torch_version}, which "
"wasn't known to have this problem.\n"
"Please open an Issue at:\n"
"https://github.com/sdatkinson/neural-amp-modeler/issues/507"
"\n"
f"and report your PyTorch version ({torch_version}) "
"so that we can keep track of versions of PyTorch that "
"might be avoided."
)
print(msg)
self._mps_65536_fallback = True
return self._forward_mps_safe(x, **kwargs)
else:
Expand Down
Loading