如何在pytorch中包含一个属性以让用户决定是否使用GPU？

如何解决如何在pytorch中包含一个属性以让用户决定是否使用GPU？

我是Python和pytorch的初学者，我为未能明确自己或不使用适当的术语而道歉。

我正在尝试编写一个线性回归类，以允许用户决定是否要在GPU上运行模型。

以下是我的代码，很愚蠢。

class LinearRegression():
    def __init__(self,dtype = torch.float64):  
        self.dtype = dtype
        self.bias = None
        self.weight = None
        
    def fit(self,x,y,std = None,device = "cuda"): 
        if x.dtype is not self.dtype:
            x = x.type(dtype = self.dtype)
        if y.dtype is not self.dtype:
            y = y.type(dtype = self.dtype)

        #let the user decide whether they want it to be standardized 
        if std == True:
          mean = torch.mean(input = x,dim=0).to(device)
          sd = torch.std(input = x,dim=0).to(device)
          x = x.to(device)
          x = (x-mean)/sd
        u = torch.ones(size = (x.size()[0],1),dtype = self.dtype).to(device)
        x_design = torch.cat([u,x],dim = 1).to(device)
        y = y.to(device)
        parameter = torch.inverse(
            torch.transpose(x_design,dim0=0,dim1=1) @ x_design).to(device) @ \
                    torch.transpose(x_design,dim1=1).to(device) @ y

        self.bias = parameter[0,0]
        self.weight = parameter[1:,0]
        return self

基本上，我只是将.to（device）添加到每个张量中，以便它可以在用户希望使用的设备上运行。 但是，我确定有更好的方法可以做到这一点，也许可以在__int__中包含一些内容？

我不确定如何使它写得更有效，这样我每次添加新功能时都不必包括.to（device）。