Ax: Crash after setting PyTorch default tensor type to cuda.FloatTensor

[ebonetti]

Hello folks at Facebook,
setting PyTorch default tensor type to cuda.FloatTensor leads Ax to crashing in a Colab GPU instance.

You can checkout a complete test setup in Colab. This is the culprit:
torch.set_default_tensor_type(torch.cuda.FloatTensor)

Keep up the great work,
Enrico Bonetti Vieno

[2timesjay]

@ebonetti thanks for bringing this to our attention. I can confirm that GPUs should be fully suppported, so we're investigating this bug now.

[Balandat]

This appears to be an issue with pytorch:

import torch
from torch.quasirandom import SobolEngine

torch.set_default_tensor_type(torch.cuda.FloatTensor)
se = SobolEngine(3)  # crashes here

We'll have to fix this upstream.

[Balandat]

See https://github.com/pytorch/pytorch/issues/32494 for the issue and https://github.com/pytorch/pytorch/pull/32496 for the fix.

[Balandat]

I landed the fix upstream in pytorch. In order for this to work you'll have to use the pytorch nightlies (starting tomorrow)

[ebonetti]

Thank you for the fast resolution, great job! 💪