Yes. I am performing the same initialization for both cpu and gpu. Here’s the architecture of the model -
BiLinear(
  (input_encoder): BiLSTM_Attention(
    (dropout): Dropout(p=0.5, inplace=False)
    (feature_embedding): Embedding(600000, 50, padding_idx=0)
    (embeddings): Embedding(
      (embeddings): Embedding(37119, 300)
    )
    (left_bilstm): BiLSTM(
      (lstm): LSTM(300, 100, bidirectional=True)
    )
    (right_bilstm): BiLSTM(
      (lstm): LSTM(300, 100, bidirectional=True)
    )
    (attention): AttentionMLP(
      (linear_1): Linear(in_features=200, out_features=100, bias=True)
      (linear_2): Linear(in_features=100, out_features=1, bias=True)
      (tanh): Tanh()
      (softmax): Softmax(dim=None)
    )
  )
  (label_encoder): LabelEncoder(
    (label_encoder): AverageLabelEncoder(
      (embeddings): Embedding(99, 300)
      (label_layer): Linear(in_features=300, out_features=20, bias=False)
    )
  )
  (input_layer): Linear(in_features=550, out_features=20, bias=False)
)

and the code that I use to train the model is:
    model = BiLinear(config).to(config.device)

    optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
    loss_fn = nn.MultiLabelMarginLoss()
    all_labels = list(config.label_to_idx.values())

    for epoch in range(config.epochs):
        print("Epoch No.", epoch+1)

        total_loss = torch.Tensor([0.]).to(config.device)

        model.train()
        idx = 0
        for x in tqdm(training_data):
            optimizer.zero_grad()
            loss = Variable(torch.Tensor([0.]))
            score_main = model(x['mention_tokens'], x['left_tokens'], x['right_tokens'], x)

            loss = loss_fn(score_main, x['labels'])
            total_loss += loss.item()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            optimizer.step()
            idx += 1

        print("Total Loss = ", total_loss)

After 1 epoch, the train loss for cpu is somewhere around 200 and the train loss for gpu is around 50492384. I’m using the train loss as a sanity check.