Dynamically choose the best device

I think it will be best if the code picked the best device automatically. This will also support Mac's mps device.

```python
def get_best_device():
  if torch.backends.mps.is_available():
    return "mps"
  elif torch.cuda.is_available():
    return "cuda"
  else:
    return "cpu"
```