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[&quot;input_shape&quot;], 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[&quot;input_shape&quot;]
        )

    def forward(self, features):
        #print(&quot;in forward&quot;)
        #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(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;)

    # 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(&quot;AE, epoch : {}/{}, loss = {:.6f}&quot;.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:

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

Thanks in advance

I see that Your model is moved to device which is decided by this line device = torch.device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;) This can be is either cpu or cuda.
So adding this line batch_features = batch_features.to(device) will actually move your input data to device.
Since your model is moved to device , You should also move your input to the device.
Below code has that change
for epoch in range(epochs_AE):
    loss = 0
    
    for batch_features in train_loader:
        batch_features = batch_features.to(device)  #this will move inout to your device        
        optimizer.zero_grad()
        outputs = model_AE(batch_features)
        train_loss = criterion(outputs, batch_features)
        ...

coming to your question : Calling .to(device) can directly move the tensor to your specified device
And if you want it it to be hard coded then do .to('cpu') or .to('cuda') on your torch tensor
