Skip to content

Mps support#120

Open
BenRogersNewsome wants to merge 2 commits into
timsainb:masterfrom
BenRogersNewsome:mps-support
Open

Mps support#120
BenRogersNewsome wants to merge 2 commits into
timsainb:masterfrom
BenRogersNewsome:mps-support

Conversation

@BenRogersNewsome

@BenRogersNewsome BenRogersNewsome commented Jan 24, 2025

Copy link
Copy Markdown

Added support for MPS graphics when using torchgate. Changes:

  • Set PYTORCH_ENABLE_MPS_FALLBACK=1 when using MPS so that 'aten::unfold_backward' doesn't fail (see General MPS op coverage tracking issue pytorch/pytorch#77764)
  • Set the dtype used in SpectralGate._read_chunk to the same datatype as the input data. This is needed since the chunks were always getting set to float64 which isn't supported on MPS. This way, any data type can be used.

@BenRogersNewsome BenRogersNewsome force-pushed the mps-support branch 3 times, most recently from ed66e9d to 73e7028 Compare January 24, 2025 14:35
This is needed since the operator 'aten::unfold_backward' is not currently implemented for MPS in pytorch - see pytorch/pytorch#77764
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant