:param channels: The last dimension of the tensor you want to apply pos emb to.
(self, channels)
| 32 | |
| 33 | class PositionalEncoding3D(nn.Module): |
| 34 | def __init__(self, channels): |
| 35 | """ |
| 36 | :param channels: The last dimension of the tensor you want to apply pos emb to. |
| 37 | """ |
| 38 | super(PositionalEncoding3D, self).__init__() |
| 39 | self.org_channels = channels |
| 40 | channels = int(np.ceil(channels / 6) * 2) |
| 41 | if channels % 2: |
| 42 | channels += 1 |
| 43 | self.channels = channels |
| 44 | inv_freq = 1.0 / (10000 ** (torch.arange(0, channels, 2).float() / channels)) |
| 45 | self.register_buffer("inv_freq", inv_freq) |
| 46 | self.cached_penc = None |
| 47 | |
| 48 | def forward(self, tensor): |
| 49 | """ |