(self, x, caches)
| 531 | return x |
| 532 | |
| 533 | def step(self, x, caches): |
| 534 | # x : (B, L, D) |
| 535 | # caches : [cache(layer) for all layers], cache : (h, inputs) |
| 536 | |
| 537 | # y : (B, L, D) |
| 538 | # caches : [cache(layer) for all layers], cache : (h, inputs) |
| 539 | |
| 540 | for i, layer in enumerate(self.layers): |
| 541 | x, caches[i] = layer.step(x, caches[i]) |
| 542 | |
| 543 | return x, caches |
| 544 | |
| 545 | |
| 546 | class BitMamba(nn.Module): |