Problem with self_supervised mode training

[zqxyus]

When i run the main.py with self_supervised mode training

the following error occurs.

Input tensor shape: torch.Size([128, 8, 100]). Additional info: {'h': 3}.
Shape mismatch, can't divide axis of length 100 in chunks of 3

The error comes from the source code in line 57-58, as is shown as below

43 class Attention(nn.Module):
44 def init(self, dim, heads=8, dropout=0.):
45 super().init()
46 self.heads = heads
47 self.scale = dim ** -0.5
48
49 self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
50 self.to_out = nn.Sequential(
51 nn.Linear(dim, dim),
52 nn.Dropout(dropout)
53 )
54
55 def forward(self, x, mask=None):
56 b, n, _, h = *x.shape, self.heads
57 qkv = self.to_qkv(x).chunk(3, dim=-1)
58 q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv)

how to solve it? Thanks

[Emadeldeen Eldele]

Try using a dividable number... try h=5