Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Move weight init of pos_embed layer to after init of other layers #4

Open
wants to merge 1 commit into
base: ai2
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Move weight init for pos_embed layer
  • Loading branch information
oliverwm1 committed Jan 30, 2024
commit f3b3c2524c9ba1ab6f4482a1ccb92d1f7daaa86a
5 changes: 4 additions & 1 deletion modulus/models/sfno/sfnonet.py
Original file line number Diff line number Diff line change
@@ -731,10 +731,13 @@ def __init__(
)
# self.pos_embed = nn.Parameter( torch.zeros(1, self.embed_dim, self.img_shape_eff[0], self.img_shape_eff[1]) )
self.pos_embed.is_shared_mp = ["matmul"]
trunc_normal_(self.pos_embed, std=0.02)

self.apply(self._init_weights)

# doing weight init of pos_embed after other layers resolves a segfault
if isinstance(self.pos_embed, nn.Parameter):
trunc_normal_(self.pos_embed, std=0.02)

def _init_weights(self, m):
"""Helper routine for weight initialization"""
if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):