try sep vit

models.py

@@ -558,6 +558,8 @@ class FouriER(torch.nn.Module):
         z = self.forward_embeddings(y)
         z = self.forward_tokens(z)
         z = z.mean([-2, -1])
+        if np.count_nonzero(np.isnan(z)) > 0:
+            print("ZZZ")
         z = self.norm(z)
         x = self.head(z)
         x = self.hidden_drop(x)
@@ -1131,6 +1133,9 @@ class PoolFormerBlock(nn.Module):
         else:
             x = x + self.drop_path(x_attn)
             x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        if np.count_nonzero(np.isnan(x)) > 0:
+            print("PFBlock")
         return x
 class PatchEmbed(nn.Module):
     """