Method set_grad_checkpointing

model/transformer_utils.py:412–413

(self, set_checkpointing=True)

        self.layernorm2 = LayerNorm(width)
    
    def set_grad_checkpointing(self, set_checkpointing=True):
        self.transformer.set_grad_checkpointing(set_checkpointing)
    
    def forward(self, x, condition=None):
        # x [B, V*N, D]