some models only return embeddings with some kwarg on forward

Author: lucidrains

setup.py

@@ -6,7 +6,7 @@
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '1.11.2',
+  version = '1.11.3',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   long_description = long_description,

vit_pytorch/accept_video_wrapper.py

@@ -43,7 +43,8 @@ def __init__(
     def forward(
         self,
         video, # (b c t h w)
-        eval_with_no_grad = False
+        eval_with_no_grad = False,
+        forward_kwargs = dict()
     ):
         add_time_pos_emb = self.add_time_pos_emb
         time = video.shape[2]
@@ -67,7 +68,7 @@ def forward(
         context = torch.no_grad if eval_with_no_grad else nullcontext
 
         with context():
-            outputs = func(video)
+            outputs = func(video, **forward_kwargs)
 
         # handle multiple outputs, say logits and embeddings returned from extractor - also handle some reduce aux loss being returned