How can I enable pytorch GPU support in Google Colab?

Question

How can I enable pytorch to work on GPU?

I've installed pytorch successfully in google colab notebook: Tensorflow reports GPU to be in place:

But torch.device function fails somehow:

How can I fix this?

Answer 1

You can enable GPU by clicking on "Change Runtime Type" under the "Runtime" menu. There is also "TPU" support available in these days.

You can define define device using torch.device:

import torch

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")