ViT CLS Token Shape

Problem

Return the tensor shape of the CLS token inside a ViT sequence.

After patch embedding the CLS token is one token per batch element with the hidden dimension. Return [batch, 1, dim].

case 1

in: [8,768]

expected: [8,1,768]

case 2

in: [1,384]

expected: [1,1,384]

case 3

in: [2,128]

expected: [2,1,128]

Starter codePython

def cls_token_shape(batch, dim):
    # Your code here
    pass

3 test cases · run your code against them in the editor