Update transformer.py

transformer.py

@@ -138,6 +138,6 @@ class ViT(nn.Module):
 
         x = self.encoder(x)
 
-        z = self.output(x, return_z=return_z)
+        output = self.output(x, return_z=return_z)
 
-        return z
+        return output