Does not use GPU support on Colab

[SK124]

The model is not at all utilising the GPU support Google colab offers. Any changes I should make to the code and in which line would be helpful.
@atong01

[atong01]

By default it is set up to use a gpu if available. Please check that one is available with torch.cuda.isavailable(). There is a --gpu flag for specifying gpu and a --use_cpu for the obvious.

[SK124]

Thanks I will look into that btw, i noticed a comment which specified that few operations might require torch<1.5 is that resolved?

[atong01]

Things were broken in pytorch 1.4 specifically for the density regularization on GPU anything 1.5 and up should be fine.

[SK124]

Is there any way i can start training from a checkpoint rather than start from scratch. which line of code should I change?

[atong01]

you should be able to load weights from a checkpoint like this, maybe inserting this around line 450 in main.py (and obviously with a correct path) would do the trick

state_dict = torch.load(args.save + "/checkpt.pth", map_location=device)
model.load_state_dict(state_dict["state_dict"])