Update README.md

README.md

@@ -34,6 +34,7 @@ import torch
 import torch.nn as nn
 import os
 from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
+from safetensors.torch import load_file 
 
 class CustomConvNeXtConfig(PretrainedConfig):
     model_type = "custom-convnext"
@@ -85,7 +86,6 @@ class CustomConvNeXtModel(PreTrainedModel):
         model = cls(config=config, model_name=model_name, num_classes=config.num_labels)
     
         # Load state_dict from safetensors file
-        from safetensors.torch import load_file  # Safetensors library
         state_dict = load_file(model_path)
         model.load_state_dict(state_dict)