(self)
| 121 | self.init_weights() |
| 122 | |
| 123 | def init_weights(self): |
| 124 | # Initialize proj_in and proj_out with trunc_normal_ |
| 125 | for m in [self.proj_in, self.proj_out]: |
| 126 | if isinstance(m, nn.Conv2d): |
| 127 | torch.nn.init.trunc_normal_(m.weight, std=0.02) |
| 128 | if m.bias is not None: |
| 129 | nn.init.zeros_(m.bias) |
| 130 | # Initialize transformer blocks |
| 131 | named_apply(init_weights_vit, self) |
| 132 | # RoPE does not need init here, it's handled internally |
| 133 | |
| 134 | def forward(self, x: Tensor) -> Tensor: |
| 135 | B, _, H, W = x.shape |