why ? input.unsqueeze(2).repeat(1, 1, nclass).transpose(1, 2)

[cheer00]

input_rep = input.unsqueeze(2).repeat(1, 1, nclass).transpose(1, 2)
task_emb = task_emb.repeat(size[0], 1)

why not use a batch (32 images ) but repeat to batch* nclass (for example 32 * 21 classes )?