Move model to device

[dariocazzani]

What does this PR do?

It moves the model to the same device as all the tensors

[ayeganov]

Maybe we should fix it in load_model?

[dariocazzani]

Unless you want to change the signature, we can't.
We need to assign model = model.to(DEVICE)

So.. we either just add that line, or we change

def load_model(model_path: str, model: nn.Module, device: torch.device) -> None

to

def load_model(model_path: str, model: nn.Module, device: torch.device) -> nn.Module
    ...
    return model

But it's messy because we are doing:

    else:
        print("No model path exists, proceeding with a new model")

In my experience, I always have a line like

model = model.to(DEVICE) 

after creation

[ayeganov]

Gotcha - then I think we can remove model.to from load_model and add this line.

[ayeganov]

But also - we can do return model.to(DEVICE) from load_model. Since we always provide a valid model pointer.

[dariocazzani]

Done.. if you disagree, I don't care

[ayeganov]

Ship it!