Update README.md

README.md

@@ -9,28 +9,23 @@ pipeline_tag: image-segmentation
 
 ## Available checkpoints
 
-We present the full OCULARNet model, a base UNet encoder - (256, 128, 64, 32) layer size - coupled with a RepVGG-B3 decoder loaded from `segmentation_models_pytorch`; and a lightweight ensemble version, OCULARNet-nano (base UNet + RepVGG-A0), for retinal artery-vein segmentation.
+We present the OCULARNet model, a base UNet encoder - (256, 128, 64, 32) layer size - coupled with a RepVGG-B3 decoder loaded from `segmentation_models_pytorch` for retinal artery-vein segmentation.
 
 Datasets, training and inference details can be found in our [GitHub repository](`https://github.com/GonzaloPlaaza/OCULAR`).
 
-
-### OCULARNet
-- OCULARNet.pth
-
-### Nano OCULARNet
-- nano_fold1.pth
-- nano_fold2.pth
-- nano_fold3.pth
-- nano_fold4.pth
-- nano_fold5.pth
-
 ## Usage
 
 ```python
 import torch
 
 device = torch.device('cuda:0' if use_cuda else 'cpu')
-model.load_state_dict(
-    torch.load("OCULARNet.pth", map_location="cpu")
-)
+checkpoint = torch.load("OCULARNet.pth", map_location=device)
+        
+if "model_state_dict" in checkpoint:
+    state_dict = checkpoint["model_state_dict"]
+else:
+    state_dict = checkpoint
+
+model.load_state_dict(state_dict)
+