(self, image: torch.Tensor, mod: torch.Tensor = None)
| 51 | |
| 52 | # @torch.compile |
| 53 | def forward(self, image: torch.Tensor, mod: torch.Tensor = None): |
| 54 | # image: [N, C, H, W] |
| 55 | # mod: [N, D] or None |
| 56 | # RGB image with [0,1] scale and properly sized |
| 57 | if self.modulation_dim is None: |
| 58 | assert mod is None, "Unexpected modulation input in dinov2 forward." |
| 59 | outs = self.model(image, is_training=True) |
| 60 | else: |
| 61 | assert mod is not None, "Modulation input is required in modulated dinov2 forward." |
| 62 | outs = self.model(image, mod=mod, is_training=True) |
| 63 | # ret = torch.cat([ |
| 64 | # outs["x_norm_clstoken"].unsqueeze(dim=1), |
| 65 | # outs["x_norm_patchtokens"], |
| 66 | # ], dim=1) |
| 67 | # return ret |
| 68 | return outs["x_norm_clstoken"].unsqueeze(dim=1), outs["x_norm_patchtokens"] |
nothing calls this directly
no outgoing calls
no test coverage detected