def set_multihead_attention(self, train: bool) -> None:
        """
        Set all multi-head attention layers into train or eval mode.

        :param train: False for evaluation mode.
        """
        import torch

        self._set_layer(train=train, layerinfo=[torch.nn.modules.MultiheadAttention])  # type: ignore

        Set all multi-head attention layers into train or eval mode.

        :param train: False for evaluation mode.
        
art/estimators/pytorch.py