Tensor for argument #2 'mat1' is on CPU, but expected it to be on GPU

Following my previous question , I have written this code to train an autoencoder and then extract the features. (There might be some changes in the variable names)

# Autoencoder class
#https://medium.com/pytorch/implementing-an-autoencoder-in-pytorch-19baa22647d1
class AE_class(nn.Module):
    def __init__(self, **kwargs):
        super().__init__()
        self.encoder_hidden_layer = nn.Linear(
            in_features=kwargs["input_shape"], out_features=128
        )
        self.encoder_output_layer = nn.Linear(
            in_features=128, out_features=128
        )
        self.decoder_hidden_layer = nn.Linear(
            in_features=128, out_features=128
        )
        self.decoder_output_layer = nn.Linear(
            in_features=128, out_features=kwargs["input_shape"]
        )

    def forward(self, features):
        #print("in forward")
        #print(type(features))
        activation = self.encoder_hidden_layer(features)
        activation = torch.relu(activation)
        code = self.encoder_output_layer(activation)
        code = torch.relu(code)
        activation = self.decoder_hidden_layer(code)
        activation = torch.relu(activation)
        activation = self.decoder_output_layer(activation)
        reconstructed = torch.relu(activation)
        return reconstructed
    
    def encode(self, features_h):
        activation_h = self.encoder_hidden_layer(features_h)
        activation_h = torch.relu(activation_h)
        code_h = self.encoder_output_layer(activation_h)
        code_h = torch.relu(code_h)
        return code_h
   

And then, for training:

def retrieve_AE_features(X_before, n_voxel_region):
    
        #  use gpu if available
    #https://discuss.pytorch.org/t/runtimeerror-tensor-for-out-is-on-cpu-tensor-for-argument-1-self-is-on-cpu-but-expected-them-to-be-on-gpu-while-checking-arguments-for-addmm/105453
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # create a model from `AE` autoencoder class
    # load it to the specified device, either gpu or cpu
    model_AE = AE_class(input_shape=n_voxel_region).to(device)

    # create an optimizer object
    # Adam optimizer with learning rate 1e-3
    optimizer = optim.Adam(model_AE.parameters(), lr=1e-3)

    # mean-squared error loss
    criterion = nn.MSELoss()

    
    
    
    X_tensor = torch.tensor(X_before, dtype=torch.float32)
    
    print(type(X_tensor))
                            
    
    train_loader = torch.utils.data.DataLoader(
        X_tensor, batch_size=64, shuffle=True, num_workers=2, pin_memory=True
    )
    
    test_loader = torch.utils.data.DataLoader(
        X_tensor, batch_size=32, shuffle=False, num_workers=2
    )
    
    print(type(train_loader))

    for epoch in range(epochs_AE):
        loss = 0
        
        for batch_features in train_loader:
            # reshape mini-batch data to [N, 784] matrix
            # load it to the active device
            #batch_features = batch_features.view(-1, 784).to(device)
            
            #print(batch_features.shape)

            # reset the gradients back to zero
            # PyTorch accumulates gradients on subsequent backward passes
            optimizer.zero_grad()

            # compute reconstructions
            outputs = model_AE(batch_features)

            # compute training reconstruction loss
            train_loss = criterion(outputs, batch_features)

            # compute accumulated gradients
            train_loss.backward()

            # perform parameter update based on current gradients
            optimizer.step()

            # add the mini-batch training loss to epoch loss
            loss += train_loss.item()

        # compute the epoch training loss
        loss = loss / len(train_loader)

        # display the epoch training loss
        
        print("AE, epoch : {}/{}, loss = {:.6f}".format(epoch + 1, epochs_AE, loss))
        
        
    #After training
    hidden_features = model_AE.encode(X_before)
    return hidden_features

However, I received the following error:

Tensor for argument #2 'mat1' is on CPU, but expected it to be on GPU (while checking arguments for addmm)

It seems some of my variables should be defined in another way to be able to be executed on GPU.

My questions:

1. How can I understand which variables will be executed on GPU and which ones on CPU?
2. How to fix it? In other words, how to define a variable executable on GPU?

Thanks in advance