Update README.md

README.md

@@ -70,16 +70,15 @@ weights_path = hf_hub_download(
         filename="model.safetensors"
     )
 
-from safetensors.torch import load_file
+from safetensors.torch import load_model
 
 model = EfficientNetSkinLesionClassifier()
-state = load_file(weights_path)
-model.load_state_dict(state, strict=True)
-
-model.eval()  # Set model to evaluation mode
+load_model(model, filename=weights_path, strict=True)
 
 model.to(device) # Don't forget to put on GPU
 
+model.eval()  # Set model to evaluation mode
+
 
 # Example with PH2 Dataset