mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
Fix mtia_extension.cpp setDevice() to correctly set current_device (#149398)
We referred to this code and found that there was a minor bug. Fix for future reference for others. Pull Request resolved: https://github.com/pytorch/pytorch/pull/149398 Approved by: https://github.com/janeyx99
This commit is contained in:
committed by
PyTorch MergeBot
parent
4f14224dc8
commit
b99e0c5412
@ -38,9 +38,8 @@ struct MTIAGuardImpl final : public c10::impl::DeviceGuardImplInterface {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void setDevice(c10::Device d) const override {
|
void setDevice(c10::Device d) const override {
|
||||||
c10::Device current_device = getDevice();
|
if (getDevice().index() != d.index()) {
|
||||||
if (current_device.index() != d.index()) {
|
current_device = d.index();
|
||||||
current_device = d;
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
void uncheckedSetDevice(c10::Device d) const noexcept override {
|
void uncheckedSetDevice(c10::Device d) const noexcept override {
|
||||||
|
|||||||
Reference in New Issue
Block a user