try sep vit

This commit is contained in:
thanhvc3 2024-04-28 11:04:26 +07:00
parent ab5c1d0b4b
commit 9075b53be6

View File

@ -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):
"""