TransformerAE throws error for different device

[ab93]

When trying to run from local using a different device than cpu, we get the following error:

Ran 1 test in 0.375s
FAILED (errors=1)
ErrorTraceback (most recent call last):
  File "/Users/abasu/Projects/Opensrc/numalogic/tests/models/autoencoder/variants/test_transformers.py", line 48, in test_transformer
    trainer.fit(model, datamodule=datamodule)
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 520, in fit
    call._call_and_handle_interrupt(
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 44, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 559, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 935, in _run
    results = self._run_stage()
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 976, in _run_stage
    self._run_sanity_check()
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1005, in _run_sanity_check
    val_loop.run()
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/utilities.py", line 177, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 115, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx)
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 375, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_kwargs.values())
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 288, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 378, in validation_step
    return self.model.validation_step(*args, **kwargs)
  File "/Users/abasu/Projects/Opensrc/numalogic/numalogic/models/autoencoder/base.py", line 92, in validation_step
    loss = self._get_reconstruction_loss(batch)
  File "/Users/abasu/Projects/Opensrc/numalogic/numalogic/models/autoencoder/variants/transformer.py", line 367, in _get_reconstruction_loss
    _, recon = self.forward(batch)
  File "/Users/abasu/Projects/Opensrc/numalogic/numalogic/models/autoencoder/variants/transformer.py", line 362, in forward
    encoded = self.encoder(batch)
  File "/Users/abasu/Library/Caches/pypoetry/virtualenvs/numalogic--00jt1TH-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/abasu/Projects/Opensrc/numalogic/numalogic/models/autoencoder/variants/transformer.py", line 212, in forward
    src = torch.add(src, _positional_encoding(feature, seq_len, device=self.device))
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, mps:0 and cpu!