CaicedoLab
/

MorphEm

Image Classification

Model card Files Files and versions

Update README.md

#1

by gouldj2 - opened 7 days ago

base: refs/heads/main

←

from: refs/pr/1

Discussion Files changed

Files changed (1) hide show

README.md +1 -1

README.md CHANGED Viewed

@@ -102,7 +102,7 @@ class PerImageNormalize(nn.Module):
 # Load model
-device = "cuda"
 model = AutoModel.from_pretrained("CaicedoLab/MorphEm", trust_remote_code=True)
 model.to(device).eval()

 # Load model
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model = AutoModel.from_pretrained("CaicedoLab/MorphEm", trust_remote_code=True)
 model.to(device).eval()