Matching requirements in main and fixing bug.

This commit to fix the RuntimeError that shows in HuggingFace through a patch found in a GitHub issue on the PyTorch Vision repository and matching some of the requirements with local.

model.py

@@ -3,10 +3,29 @@ import torch
 from pathlib import Path
 from torch import nn
 from torchvision.models import efficientnet_b2, EfficientNet_B2_Weights
+from torchvision.models._api import WeightsEnum
+from torch.hub import load_state_dict_from_url
 from typing import Optional, Tuple
 
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
+# A RuntimeError appeared in HuggingFace when the application tried to load the
+# weights of the model. The following link is to the source of the fix now
+# implemented.
+# https://github.com/pytorch/vision/issues/7744#issuecomment-1757321451
+
+
+def get_state_dict(self, *args, **kwargs):
+    """
+        Override intented to fix a bug while loading the state_dict
+        from the internet.
+    """
+    kwargs.pop("check_hash")
+    return load_state_dict_from_url(self.url, *args, **kwargs)
+
+
+WeightsEnum.get_state_dict = get_state_dict
+
 
 def create_effnetb2_model(
     num_classes: int,
@@ -18,7 +37,7 @@ def create_effnetb2_model(
     weights = EfficientNet_B2_Weights.DEFAULT
     transforms = weights.transforms()
     model = efficientnet_b2(weights=weights)
-    
+
     model.classifier = nn.Sequential(
       nn.Dropout(p=0.3, inplace=True),
       nn.Linear(in_features=1408, out_features=num_classes, bias=True)
@@ -28,5 +47,5 @@ def create_effnetb2_model(
         model.load_state_dict(torch.load(st_dict, map_location=DEVICE))
     for param in model.parameters():
         param.requires_grad = False
-    
+
     return model.to(DEVICE), transforms

requirements.txt

@@ -1,3 +1,3 @@
-torch>=2.0.0
-torchvision>=0.15.0
-gradio>=3.33.1
+torch>=2.1.0
+torchvision>=0.16.0
+gradio==3.33.1