⚡️ Enable training with CUDA #59
Conversation
|
Looks like |
There was a problem hiding this comment.
Thank you for your PR. Unfortunately, I think there is a better approach to address #58.
The cast to .int() in MaskedMLP is there because matmul is not possible for boolean tensors on CPU (and CUDA). We can simply cast to .double() instead to support CUDA. I tried in b5fc6cc and it works perfectly.
Concerning the CI, instead of modifying many (but not all) tests, it would be better to add a global option to pytest to run all tests with set_device_default("cuda"). I tried in 2f325e4, which allows to run
pytest --device cuda
| # PyTorch doesn't support this operation for integer arrays on CUDA devices | ||
| precedence = ( | ||
| adjacency.cpu().int() @ adjacency.cpu().int().t() == adjacency.sum(dim=-1).cpu() | ||
| ) | ||
| try: | ||
| precedence = precedence.to(torch.get_default_device()) | ||
| except AttributeError: | ||
| precedence = torch.tensor(precedence.detach().cpu().numpy()) |
There was a problem hiding this comment.
The cast to .int() is here because matmul is not possible for boolean tensors on CPU (and CUDA). We can simply cast to .double() instead to support CUDA.
|
No longer needed -- #58 has been fixed! |
PR implementing the feature suggested in #58. The integer matmul operation needed for a masked autoregressive transform should be explicitly carried out on CPU and then migrated to the current default device. All other tensors get created on the default device anyway. This way, if the user calls
torch.set_default_device('cuda'), that preference will be followed.I have modified existing tests to make sure that no existing code should be broken by my modification; all tests are now performed on CPU and with CUDA to be sure they pass either way.