Upload model

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1afa15c915cfba02cf65f9fda685e845a333113fdf807050238d1d8136c8a7ec
+oid sha256:997619d28e6a8e5a1c7ceaab745e1593fa761e00bd59afb6086c3aa480975bf6
 size 247728

modeling_lenet.py

@@ -27,8 +27,8 @@ class LeNet(torch.nn.Module):
             )
         )
         
-    def forward(self, tensor) -> torch.Tensor:
-        return self.model(tensor)
+    def forward(self, pixel_values) -> torch.Tensor:
+        return self.model(pixel_values)
         
 
 class LeNetModel(PreTrainedModel):
@@ -38,8 +38,8 @@ class LeNetModel(PreTrainedModel):
         super().__init__(config)
         self.model = LeNet()
 
-    def forward(self, tensor):
-        return self.model.forward_features(tensor)
+    def forward(self, pixel_values):
+        return self.model.forward_features(pixel_values)
         
 
 class LeNetModelForImageClassification(PreTrainedModel):
@@ -49,8 +49,8 @@ class LeNetModelForImageClassification(PreTrainedModel):
         super().__init__(config)
         self.model = LeNet()
 
-    def forward(self, tensor, labels=None):
-        logits = self.model(tensor)
+    def forward(self, pixel_values, labels=None):
+        logits = self.model(pixel_values)
         if labels is not None:
             loss = torch.nn.functional.cross_entropy(logits, labels)
             return {"loss": loss, "logits": logits}