simplified setup.py

2021-01-30 02:44:22 +09:00
parent a19dcf967b
commit 77efe70a0a
8 changed files with 57 additions and 205 deletions
--- a/clip/clip.py
+++ b/clip/clip.py
@ -75,6 +75,8 @@ def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_a

    if not jit:
        model = build_model(model.state_dict()).to(device)
+        if device == "cpu":
+            model.float()
        return model, transform

    # patch the device names