From 10817696fb3e06fa6eca6d7b2dbbe24d2b2ca912 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Tue, 24 Sep 2024 10:34:11 -0600
Subject: [PATCH] Fixed issue where direct vision was not passing additional
 params from resampler when it is added

---
 toolkit/custom_adapter.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/toolkit/custom_adapter.py b/toolkit/custom_adapter.py
index e594378b..23367bbf 100644
--- a/toolkit/custom_adapter.py
+++ b/toolkit/custom_adapter.py
@@ -889,6 +889,8 @@ class CustomAdapter(torch.nn.Module):
                     yield from attn_processor.parameters(recurse)
                 if self.config.train_image_encoder:
                     yield from self.vision_encoder.parameters(recurse)
+                if self.config.num_tokens:
+                    yield from self.vd_adapter.resampler.parameters(recurse)
         elif self.config.type == 'te_augmenter':
             yield from self.te_augmenter.parameters(recurse)
             if self.config.train_image_encoder: