Update vit_fnqs_model.py

vit_fnqs_model.py

@@ -25,10 +25,10 @@ class ViTFNQSModel(FlaxPreTrainedModel):
                              transl_invariant=config.tras_inv,
                              two_dimensional=config.two_dim,
         )
-        # if not "return_z" in kwargs:
-        #     self.return_z = False
-        # else:
-        #     self.return_z = kwargs["return_z"]
+        if not "return_z" in kwargs:
+            self.return_z = False
+        else:
+            self.return_z = kwargs["return_z"]
 
         super().__init__(config, ViTFNQS, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)