Hi,
I am trying to implement a simple autoencoder class. I want to run it on GPU, and therefore have tried to explicity move everything to CUDA (using “.to(self.device)”). However, I’m still getting an “expected CPU” error. Could I get help identifying what additional parts need to be moved to cuda? The class looks like the below:
class AE(torch.nn.Module):
    def __init__(self, Xtr, Xval,
                 eps = 1e-7, weight_decay = 1e-5, C=1., lr = 1e-2, 
                 device=torch.device('cuda')):
        super(AE, self).__init__()
        self.weight_decay = weight_decay
        self.lr = lr
        self.C = C
        self.device = device
        self.__initfact__(Xtr, Xval)

    def __initfact__(self, Xtr, Xval):
        self.ntr = torch.tensor(Xtr.shape, dtype=int)[0].to(self.device)
        self.nval = torch.tensor(Xval.shape, dtype=int)[0].to(self.device)
        self.m = torch.tensor(Xtr.shape, dtype=int)[1].to(self.device)
        self.Xtr = torch.from_numpy(Xtr).float().to(self.device)
        self.Xval = torch.from_numpy(Xval).float().to(self.device)

        self.loss_rcs = nn.MSELoss().to(self.device)

        self.encoder = nn.Sequential(
            nn.Linear(self.m, 2048).to(self.device), 
            nn.ReLU().to(self.device),
            nn.Linear(2048, 1024).to(self.device)
        ).to(self.device)
        self.decoder = nn.Sequential(
            nn.Linear(1024, 2048).to(self.device),
            nn.ReLU().to(self.device),
            nn.Linear(2048, self.m).to(self.device)
        ).to(self.device)
        
        self.opt = torch.optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
    
    def forward(self, X):
        X = self.encoder(X)
        X = self.decoder(X)
        return X

When I initialize this class with Numpy arrays Xtr Xval, this is the error that I get:
raceback (most recent call last):
  File "pretrain/modeling/pretrain_ae.py", line 61, in <module>
    m = AE(
  File "/home/srd6051/pretrain/modeling/AE.py", line 33, in __init__
    self.__initfact__(Xtr, Xval)
  File "/home/srd6051/pretrain/modeling/AE.py", line 73, in __initfact__
    nn.Linear(2048, self.m).to(self.device)
  File "/home/srd6051/anaconda3/envs/bbcarenv/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 81, in __init__
    self.bias = Parameter(torch.Tensor(out_features))
TypeError: expected CPU (got CUDA)

I am using pytorch version 1.8.1. I’m happy to provide any additional information that might be useful. Any insight would be greatly appreciated! Thank you in advance.