| 105 | self.diffusion_batch_mul = diffusion_batch_mul |
| 106 | |
| 107 | def initialize_weights(self): |
| 108 | # parameters |
| 109 | torch.nn.init.normal_(self.class_emb.weight, std=.02) |
| 110 | torch.nn.init.normal_(self.fake_latent, std=.02) |
| 111 | torch.nn.init.normal_(self.mask_token, std=.02) |
| 112 | torch.nn.init.normal_(self.encoder_pos_embed_learned, std=.02) |
| 113 | torch.nn.init.normal_(self.decoder_pos_embed_learned, std=.02) |
| 114 | torch.nn.init.normal_(self.diffusion_pos_embed_learned, std=.02) |
| 115 | |
| 116 | # initialize nn.Linear and nn.LayerNorm |
| 117 | self.apply(self._init_weights) |
| 118 | |
| 119 | def _init_weights(self, m): |
| 120 | if isinstance(m, nn.Linear): |