Change comment: incorrect shape

[stefanvasilev]

dualprompt-pytorch/prompt.py

Line 118 in 7eb457d

batched_prompt_raw = self.prompt[:,:,idx] # num_layers, B, top_k, length, C

The comment here suggests the wrong shape # num_layers, B, top_k, length, C, should be # num_layers, dual, B, top_k, length, num_heads, embed_dim // num_heads

[tunglamlqddb]

dualprompt-pytorch/prompt.py

Line 120 in 7eb457d

batched_prompt = batched_prompt_raw.reshape(

And I guess there should be a permutation before this reshape to correct the position of batch_size:

batched_prompt_raw = batched_prompt_raw.permute(0, 2, 1, 3, 4, 5, 6)
batched_prompt = batched_prompt_raw.reshape(
num_layers, batch_size, dual, top_k * length, num_heads, heads_embed_dim
)