RuntimeError: Input tensors need to be on the same GPU

chrathans: <p>I want to train a hugging face <code>gpt2</code> model using the <code>wikitext</code> dataset. I am loading the model the tokenizer applies it to the data and finally, I am training the model using these data. I am using the package <code>transformers-heads</code> and the code in this <a href="https://github.com/center-for-humans-and-machines/transformer-heads/blob/main/notebooks/gpt2/linear_probe.ipynb" rel="noopener nofollow ugc">link</a>  which is actually the following:</p>
<pre><code class="lang-auto"># GPT2 is the fastest and requires the fewest memory. However, this works just the same with any Llama or Mistral model. Just change model_path to its hugging face path.
  model_path = "gpt2"
  train_epochs = 1
  eval_epochs = 1
  logging_steps = 100

  model_params = get_model_params(model_path)
  model_class = model_params["model_class"]
  hidden_size = model_params["hidden_size"]
  vocab_size = model_params["vocab_size"]
  print(model_params)

  model_params = get_model_params(model_path)
  model_class = model_params["model_class"]
  hidden_size = model_params["hidden_size"]
  vocab_size = model_params["vocab_size"]
  print(model_params)

  heads_configs = [
      HeadConfig(
          name="wikitext_head",
          layer_hook=-4,  # Hook to layer [-4] (Drop 3 layers from the end)
          in_size=hidden_size,
          num_layers=1,
          output_activation="linear",
          is_causal_lm=True,
          loss_fct="cross_entropy",
          num_outputs=vocab_size,
          is_regression=False,
          output_bias=False,
      )
  ]

  dd = load_dataset("wikitext", "wikitext-2-v1")
  #device
  tokenizer = AutoTokenizer.from_pretrained(model_path)

  if tokenizer.pad_token_id is None:
      tokenizer.pad_token = tokenizer.eos_token

  def tokenize_function(examples):
      out = tokenizer(examples["text"], padding=False, truncation=True)
      out[heads_configs[0].name] = out["input_ids"].copy()
      return out.to(device)

  for split in dd.keys():
      dd[split] = dd[split].filter(function=lambda example: len(example["text"]) &gt; 10)
      dd[split] = dd[split].map(tokenize_function, batched=True)

  dd.set_format(type="torch", columns=["input_ids", "attention_mask", heads_configs[0].name])
  
  for split in dd.keys():
      dd[split] = dd[split].remove_columns("text")
  
  quantization_config = BitsAndBytesConfig(
      load_in_4bit=True,
      load_in_8bit=False,
      llm_int8_threshold=6.0,
      llm_int8_has_fp16_weight=False,
      bnb_4bit_compute_dtype=torch.float32,
      bnb_4bit_use_double_quant=True,
      bnb_4bit_quant_type="nf4",
  )

  model = load_headed(
      model_class,
      model_path,
      head_configs=heads_configs,
      quantization_config=quantization_config,
      device_map=device,
  )
  print_trainable_parameters(model)

 #print(get_top_n_preds(n=5, model=model, text="The historical significance of", tokenizer=tokenizer))
  args = TrainingArguments(
      output_dir="linear_probe_test",
      learning_rate=0.0002,
      num_train_epochs=train_epochs,
      logging_steps=logging_steps,
      do_eval=False,
      remove_unused_columns=False,  # Important to set to False, otherwise things will fail
  )
  collator = DataCollatorWithPadding(
      feature_name_to_padding_value={
          "input_ids": tokenizer.pad_token_id,
          heads_configs[0].name: -100,
          "attention_mask": 0,
      }
  )
  trainer = Trainer(
      model,
      args=args,
      train_dataset=dd["train"],
      data_collator=collator
  )

  pdb.set_trace()
  trainer.train()

  print(evaluate_head_wise(model, dd["validation"], collator, epochs=eval_epochs))
  print(evaluate_head_wise(model, dd["validation"], collator, epochs=eval_epochs))
</code></pre>
<p>When running the <code>trainer.train()</code> I am receiving the following error:</p>
<p><code>RuntimeError: Input tensors need to be on the same GPU, but found the following tensor and device combinations: [(torch.Size([1, 884736]), device(type='cuda', index=1)), (torch.Size([27648]), device(type='cuda', index=0)), (torch.Size([2304, 768]), device(type='cuda', index=1))]</code></p>
<p>I guess different things are running on different GPUs. Why is that an issue? And how can I locate what is running on different GPUs?</p>