Method __init__

sat/sgm/modules/encoders/modules.py:273–286

(
        self,
        n_register_token=10,
        dim=512,
        freeze=False,
    )

    """Uses the T5 transformer encoder for text"""

    def __init__(
        self,
        n_register_token=10,
        dim=512,
        freeze=False,
    ):
        super().__init__()
        self.n_register_token = n_register_token
        self.dim = dim
        self.register_tokens = nn.Parameter(
            torch.randn(n_register_token, dim)
        )
        if freeze:
            self.freeze()

    def freeze(self):
        self.register_tokens.requires_grad = False