Computes the diagonal and the full matrix of the covariance matrix for activations from a Conv2d layer. Parameters: - h (Ten

:param input: batch_size * in_c * in_h * in_w :param grad_output: batch_size * out_c * h * w :param layer: nn.module batch_si

Args: x: input features with shape of (B, H, W, C)

Multi-headed attention. This applies the multi-headed attention on the input tensors `hidden_states`. For a single sequence (for simp

Mini GPT-1. This is a small version of OpenAI's GPT-1 transformer for (causal) language modeling. This module returns for each positi

Embedding module for GPT-1. This module combines token and positional embeddings, to return the embeddings from a sequence of input t

Return last computed learning rate by current scheduler.

Calculates the learning rate at batch index. This function treats `self.last_epoch` as the last batch index. If `self.cycle_momentum`

Computes gradient covariance diagonal for LayerNorm. Args: s: Output gradients layer: LayerNorm layer

Computes the diagonal and the full matrix of the covariance matrix for activations from a LayerNorm layer. Parameters: - h (

:param input: batch_size * input_dim :param grad_output: batch_size * output_dim :param layer: [nn.module] output_dim * input

Computes gradient covariance diagonal for Linear layer. Args: s: Output gradients layer: Linear layer

Computes the diagonal and the full matrix of the covariance matrix for activations from a Linear layer. Parameters: - h (Ten

:param input: batch_size * input_dim :param grad_output: batch_size * output_dim :param layer: [nn.module] output_dim * input

Load the Kronecker factors from the dictionary state of the optimizer Args: state_dict (Dict[str:List[Tensor]]): Dictionnary cont

Loads the schedulers state. Arguments: state_dict (dict): scheduler state. Should be an object returned from a ca

Loads the schedulers state. Arguments: state_dict (dict): scheduler state. Should be an object returned from a ca

r"""ResNeXt-101 32x8d model from `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_

r"""ResNeXt-50 32x4d model from `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_

Save the Kronecker factors in the dictionary state of the optimizer Returns: Dict[str:List[Tensor]]: Dictionnary containing the s

Returns the state of the scheduler as a :class:`dict`. It contains an entry for every variable in self.__dict__ which is not the opti

Returns the state of the scheduler as a :class:`dict`. It contains an entry for every variable in self.__dict__ which is not the opti

Performs a single optimization step. Arguments: closure (callable, optional): A closure that reevaluates the model

Performs a single optimization step across all parameter groups of the model. This method iterates over each parameter group, updatin

Performs a single optimization step across all parameter groups of the model. This method iterates over each parameter group, updatin

Performs a single optimization step. Arguments: gradsH: The gradient used to compute Hessian vector product. closure (

Performs a single optimization step. Arguments: closure (callable, optional): A closure that reevaluates the model