one_hot to cuda

[ironv]

My simple model for a binary classification problem is shown below. When I run this code using trainer = pl.Trainer(accelerator="gpu", ... I need the .cuda() at the end of one_hot or else I get the error RuntimeError: Input and parameter tensors are not at the same device, found input tensor at cpu and parameter tensor at cuda:0. But when I set accelerator=cpu, I get an error. Is there a better way of handling this?

class RNNModel(nn.Module):
  
  def __init__(self, prm):
    super().__init__()
        
    #LSTM
    self.rnn_ab = nn.LSTM(input_size=len(let2ind)+1, hidden_size=prm["lstm_hidden_size"], num_layers=1, batch_first=True)
    self.rnn_ag = nn.LSTM(input_size=len(let2ind)+1, hidden_size=prm["lstm_hidden_size"], num_layers=1, batch_first=True)
    
    #LSTM output to be concatenated
    self.fc1 = nn.Linear(in_features=2*prm["lstm_hidden_size"], out_features=16)
    self.dropout = nn.Dropout(p=.5)
    self.fc2 = nn.Linear(in_features=16, out_features=1)
  
  def forward(self, ab_tens, ag_tens):    
    #antibody
    ab_out = F.one_hot(ab_tens,len(let2ind)+1).type(torch.FloatTensor).cuda()
    ab_out,(ab_hid,_) = self.rnn_ab(ab_out)
    
    #antigen
    ag_out = F.one_hot(ag_tens,len(let2ind)+1).type(torch.FloatTensor).cuda()
    ag_out,(ag_hid,_) = self.rnn_ab(ag_out)
    
    out = self.fc1(torch.cat([ab_out[:,-1,:],ag_out[:,-1,:]],dim=1))
    out = self.dropout(F.relu(out))
    out = self.fc2(out)
    return out.squeeze()

[akihironitta]

@ironv That's because .type(torch.FloatTensor) moves the tensor to CPU, so instead, try .type(torch.float32).