NagaSaiAbhinay commited on
Commit
ff4f79e
·
verified ·
1 Parent(s): 2ffab58

Correct the forward call

Browse files
Files changed (1) hide show
  1. model.py +1 -1
model.py CHANGED
@@ -13,7 +13,7 @@ class CSDModel(PreTrainedModel):
13
 
14
  @torch.inference_mode()
15
  def forward(self, pixel_values):
16
- features = self.backbone(pixel_values)
17
  style_embeds = self.out_style(features)
18
  content_embeds = self.out_content(features)
19
  return features, style_embeds, content_embeds
 
13
 
14
  @torch.inference_mode()
15
  def forward(self, pixel_values):
16
+ features = self.backbone(pixel_values, return_dict=False)[1]
17
  style_embeds = self.out_style(features)
18
  content_embeds = self.out_content(features)
19
  return features, style_embeds, content_embeds