Here is the full trace @BramVanroy:
Input tensor is not an XLA tensor: torch.FloatTensor
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 330, in _mp_start_fn
    _start_fn(index, pf_cfg, fn, args)
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 324, in _start_fn
    fn(gindex, *args)
  File "<ipython-input-8-8b7fe8f75065>", line 286, in _mp_fn
    main()
  File "<ipython-input-8-8b7fe8f75065>", line 257, in main
    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
  File "/usr/local/lib/python3.6/dist-packages/transformers/trainer.py", line 521, in train
    tr_loss += self.training_step(model, inputs, optimizer)
  File "/usr/local/lib/python3.6/dist-packages/transformers/trainer.py", line 702, in training_step
    outputs = model(**inputs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 726, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/transformers/modeling_t5.py", line 1141, in forward
    return_tuple=return_tuple,
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 726, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/transformers/modeling_t5.py", line 703, in forward
    inputs_embeds = self.embed_tokens(input_ids)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 726, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/sparse.py", line 126, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py", line 1845, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: torch_xla/csrc/aten_xla_bridge.cpp:69 : Check failed: xtensor 
*** Begin stack trace ***
	tensorflow::CurrentStackTrace[abi:cxx11]()
	torch_xla::bridge::GetXlaTensor(at::Tensor const&)
	torch_xla::AtenXlaType::index_select(at::Tensor const&, long, at::Tensor const&)
	c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoRuntimeFunctor_<at::Tensor (*)(at::Tensor const&, long, at::Tensor const&), at::Tensor, c10::guts::typelist::typelist<at::Tensor const&, long, at::Tensor const&> >, at::Tensor (at::Tensor const&, long, at::Tensor const&)>::call(c10::OperatorKernel*, at::Tensor const&, long, at::Tensor const&)
	
	at::Tensor::index_select(long, at::Tensor const&) const
	at::native::embedding(at::Tensor const&, at::Tensor const&, long, bool, bool)
	torch_xla::AtenXlaType::embedding(at::Tensor const&, at::Tensor const&, long, bool, bool)
	c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoRuntimeFunctor_<at::Tensor (*)(at::Tensor const&, at::Tensor const&, long, bool, bool), at::Tensor, c10::guts::typelist::typelist<at::Tensor const&, at::Tensor const&, long, bool, bool> >, at::Tensor (at::Tensor const&, at::Tensor const&, long, bool, bool)>::call(c10::OperatorKernel*, at::Tensor const&, at::Tensor const&, long, bool, bool)
	
	at::embedding(at::Tensor const&, at::Tensor const&, long, bool, bool)
	
	
	
	at::embedding(at::Tensor const&, at::Tensor const&, long, bool, bool)
	
	
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	_PyFunction_FastCallDict
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	_PyFunction_FastCallDict
	
	
	_PyObject_FastCallKeywords
	
	_PyEval_EvalFrameDefault
	
	_PyFunction_FastCallDict
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	_PyFunction_FastCallDict
	
	
	_PyObject_FastCallKeywords
	
	_PyEval_EvalFrameDefault
	
	_PyFunction_FastCallDict
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	_PyFunction_FastCallDict
	
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	_PyFunction_FastCallDict
	
	
	
	_PyObject_FastCallKeywords
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	_PyFunction_FastCallDict
	
	PyObject_Call
	_PyEval_EvalFrameDefault
*** End stack trace ***
Input tensor is not an XLA tensor: torch.FloatTensor
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-10-e0997d11f847> in <module>()
      2 from numpy import inf
      3 
----> 4 xmp.spawn(_mp_fn, args=(), nprocs=8, start_method='fork')

2 frames
/usr/local/lib/python3.6/dist-packages/torch/multiprocessing/spawn.py in join(self, timeout)
    111                 raise Exception(
    112                     "process %d terminated with exit code %d" %
--> 113                     (error_index, exitcode)
    114                 )
    115 

Exception: process 1 terminated with exit code 17

(This is repeated multiple times for each TPU core so I omitted the duplicate text).
I also tried using your tip of explicitly returning the tensors using the to(xm.xla_device()) and got the same error message.