added cpu support to img2img.py

3 years ago · d3f1ebcbd7
parent 333e21ed80
commit d3f1ebcbd7
1 changed files with 2 additions and 1 deletions
--- a/scripts/img2img.py
+++ b/scripts/img2img.py
@ -40,6 +40,7 @@ def load_model_from_config(config, ckpt, verbose=False):
        print("unexpected keys:")
        print(u)
    if torch.cuda.is_available():
        model.cuda()
    model.eval()
    return model