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

EN

class inside(nn.Module):
    def __init__(self):
        super(inside, self).__init__()        
        self.weight_h = nn.Parameter(SOMETHING GOES HERE) # send to CPU or GPU
        self.weight_v = nn.Conv2d(SOMETHING GOES HERE) # send to CPU or GPU

    def forward(self, x):
        ...
        return x

class main_class(nn.Module):
    def __init__(self):
        super(main_class, self).__init__()        
        self.paramA = nn.Conv2d(...)
        self.paramB = nn.Conv2d(...)
        self.in_class = inside()

    def forward(self, x):
        ...
        return x

device = #Define what GPU device to use or CPU

object = main_class()
object = object.to(device)