pytorch/torch/mtia
Pat Vignola 769f19bf95 [MTIA] Add _mtia_exchangeDevice to MTIA module (#149322)
Summary: The FlexAttention path uses `_exchange_device`, so it will be needed eventually for MTIA as well.

Test Plan: `buck2 test fbcode//mtia/host_runtime/torch_mtia/tests:test_torch_mtia_api -- test_exchange_device`

Reviewed By: chaos5958

Differential Revision: D70072059

Pull Request resolved: https://github.com/pytorch/pytorch/pull/149322
Approved by: https://github.com/chaos5958
2025-03-17 19:31:10 +00:00
..
__init__.py [MTIA] Add _mtia_exchangeDevice to MTIA module (#149322) 2025-03-17 19:31:10 +00:00
_utils.py
memory.py [MTIA] (4/n) Implement PyTorch APIs to query/reset device peak memory usage (#146751) 2025-02-11 03:51:48 +00:00