Update model.py

model.py

@@ -1,3 +1,4 @@
+import torch
 import torch.nn as nn
 from .config import CSDConfig
 from transformers import PreTrainedModel, CLIPVisionModel
@@ -10,6 +11,7 @@ class CSDModel(PreTrainedModel):
         self.out_style = nn.Linear(config.hidden_size, config.style_projection_dim, bias=False)
         self.out_content = nn.Linear(config.hidden_size, config.content_projection_dim, bias=False)
 
+    @torch.inference_mode()
     def forward(self, pixel_values):
         features = self.backbone(pixel_values)
         style_embeds = self.out_style(features)