Update modeling_resnet.py

modeling_resnet.py

@@ -1,13 +1,14 @@
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import ImageClassifierOutput
 
 class ResNetConfig(PretrainedConfig):
     model_type = "custom_resnet"
     def __init__(self, num_classes=10, num_channels=3, **kwargs):
         super().__init__(**kwargs)
         self.num_classes = num_classes
-        self.num_channels = num_channels  
+        self.num_channels = num_channels
 
 class BasicBlock(nn.Module):
     expansion = 1
@@ -61,7 +62,8 @@ class ResNet(PreTrainedModel):
             layers.append(block(self.in_channels, out_channels))
         return nn.Sequential(*layers)
 
-    def forward(self, x, labels=None):
+    def forward(self, pixel_values=None, labels=None, **kwargs):
+        x = pixel_values
         x = self.relu(self.bn1(self.conv1(x)))
         x = self.maxpool(x)
         x = self.layer1(x)
@@ -71,4 +73,12 @@ class ResNet(PreTrainedModel):
         x = torch.flatten(self.avgpool(x), 1)
         logits = self.fc(x)
         
-        return {"logits": logits}
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+            
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits
+        )