| 35 | **kwargs: passed to `encoder.modules.transformer.StreamingTransformerEncoder`. |
| 36 | """ |
| 37 | def __init__(self, n_q: int = 32, card: int = 1024, dim: int = 200, **kwargs): |
| 38 | super().__init__() |
| 39 | self.card = card |
| 40 | self.n_q = n_q |
| 41 | self.dim = dim |
| 42 | self.transformer = m.StreamingTransformerEncoder(dim=dim, **kwargs) |
| 43 | self.emb = nn.ModuleList([nn.Embedding(card + 1, dim) for _ in range(n_q)]) |
| 44 | self.linears = nn.ModuleList([nn.Linear(dim, card) for _ in range(n_q)]) |
| 45 | |
| 46 | def forward(self, indices: torch.Tensor, |
| 47 | states: tp.Optional[tp.List[torch.Tensor]] = None, offset: int = 0): |