如何在Jupyter中设置PyTorch默认在M1 GPU上创建张量？

你可以通过修改PyTorch的全局设置，实现无需手动指定device就能让所有新创建的张量默认在MPS GPU上运行，具体方案如下：

方法1：直接设置默认设备 #

在Jupyter Notebook的起始单元格中执行以下代码：

import torch

# 将默认设备设置为MPS
torch.set_default_device('mps')

之后创建的所有张量都会自动分配到MPS设备上，示例：

a = torch.randn(())
print(a.device)  # 输出结果为 mps:0

方法2：同时指定默认设备与数据类型 #

如果需要统一默认数据类型，可以搭配set_default_dtype使用：

import torch

# 设置默认数据类型（按需调整，比如torch.float64）
torch.set_default_dtype(torch.float32)
# 设置默认设备为MPS
torch.set_default_device('mps')

后续创建张量时无需额外参数：

b = torch.randn(3, 3)
print(b.device)  # 输出 mps:0
print(b.dtype)   # 输出 torch.float32

注意事项 #

• 该设置仅对当前Notebook会话生效，重启Notebook后需要重新执行代码
• 若部分操作不支持MPS，PyTorch会抛出错误，此时可通过tensor.to('cpu')手动将张量移回CPU处理

内容的提问来源于stack exchange，提问作者Kliker