将一个类的所有参数和对象发送到PyTorch中的同一设备

import torch

class MyClass(torch.nn.Module):
    def __init__(self):
        super(MyClass, self).__init__()
        self.param1 = torch.nn.Parameter(torch.randn(3, 3))
        self.param2 = torch.nn.Parameter(torch.randn(3, 3))
        self.tensor1 = torch.randn(3, 3)
        self.tensor2 = torch.randn(3, 3)

    def to(self, device):
        self.param1 = self.param1.to(device)
        self.param2 = self.param2.to(device)
        self.tensor1 = self.tensor1.to(device)
        self.tensor2 = self.tensor2.to(device)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
my_object = MyClass()
my_object.to(device)