Method forward

astroclip/models/astroclip.py:218–235

(self, x: torch.tensor, return_weights: bool = False)

        )

    def forward(self, x: torch.tensor, return_weights: bool = False):
        # Pass through the backbone
        with torch.set_grad_enabled(not self.freeze_backbone):
            x = self.backbone.patch_embed(x)
            for blk in self.backbone.blocks:
                x = blk(x)
            embedding = self.backbone.norm(x)

        # Pass through cross-attention
        x, attentions = self.cross_attention(embedding)

        # Pass through MLP and residual connection
        x = self.mlp(x)

        if return_weights:
            return x.squeeze(), attentions[1]

        return x.squeeze()


class SpectrumHead(nn.Module):