Update resampler.py (#3)

- Update resampler.py (cc31e963f29bb5e66f0d399572ba0650e7fc92d4)


Co-authored-by: ed <[email protected]>

resampler.py

@@ -160,6 +160,15 @@ class Resampler(nn.Module):
             nn.init.constant_(m.bias, 0)
             nn.init.constant_(m.weight, 1.0)
 
+    def _initialize_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
     def forward(self, x, tgt_sizes=None, temporal_ids=None):
         assert x.shape[0] == tgt_sizes.shape[0]
         bs = x.shape[0]