Update moondream.py
Browse filesfix: one method absorbed another method.
- moondream.py +11 -11
moondream.py
CHANGED
|
@@ -171,17 +171,17 @@ class MoondreamModel(nn.Module):
|
|
| 171 |
blk.kv_cache.v_cache = torch.zeros(shape, device=device, dtype=dtype)
|
| 172 |
|
| 173 |
|
| 174 |
-
|
| 175 |
-
|
| 176 |
-
|
| 177 |
-
|
| 178 |
-
|
| 179 |
-
|
| 180 |
-
|
| 181 |
-
|
| 182 |
-
|
| 183 |
-
|
| 184 |
-
|
| 185 |
|
| 186 |
@property
|
| 187 |
def device(self):
|
|
|
|
| 171 |
blk.kv_cache.v_cache = torch.zeros(shape, device=device, dtype=dtype)
|
| 172 |
|
| 173 |
|
| 174 |
+
def _setup_caches(self):
|
| 175 |
+
c = self.config.text
|
| 176 |
+
for b in self.text.blocks:
|
| 177 |
+
b.kv_cache = KVCache(
|
| 178 |
+
c.n_heads,
|
| 179 |
+
c.n_kv_heads,
|
| 180 |
+
c.max_context,
|
| 181 |
+
c.dim,
|
| 182 |
+
device=self.device,
|
| 183 |
+
dtype=self.vision.pos_emb.dtype,
|
| 184 |
+
)
|
| 185 |
|
| 186 |
@property
|
| 187 |
def device(self):
|