Confusion in "mean_patch" pooling with ViT encoders

[ashish-221b]

In vit.py in forward_set function

the transform module is being used twice in 1 forward evaluation.

First the common forward

few-shot-diffusion-models/model/vit.py

Lines 200 to 201 in 8bba697

	x = self.dropout(x)
	x_set = self.transformer(x)

Second during mean_patch pooling

few-shot-diffusion-models/model/vit.py

Lines 220 to 225 in 8bba697

	elif self.pool == "mean_patch":
	x = x_set[:, self.k:]
	# attention here what you average
	x = x.view(b, np//ns, ns, -1)
	x = x.mean(dim = 2)
	x = self.transformer(x)

Is this intended?