//  [python]
def step(self, update_stats=True, update_params=True):
        # Performs one step of preconditioning.
        fisher_norm = 0.0
        for group in self.param_groups:
            # print(torch.cuda.memory_allocated()/(1024**3))
            # Getting parameters
            if len(group["params"]) == 2:
                weight, bias = group["params"]
            else:
                weight = group["params"][0]
                bias = None
            state = self.state[weight]

            # print(group['layer_type'])
            if "BatchNorm" in group["layer_type"] and self.use_batch_norm:
                # now compute hessian of weights
                diag_comp = (
                    100
                    * weight.size(0)
                    * self.eps
                    * torch.eye(
                        weight.size(0), device=weight.device, dtype=weight.dtype
                    )
                )
                # print(weight.size())
                weight_hessian = jacobian(weight.grad, weight) + diag_comp
                # print(weight_hessian)

                weight_inv_chol = torch.linalg.cholesky(weight_hessian)
                state["w_ic"] = weight_inv_chol

                if bias is not None:
                    diag_comp = (
                        100
                        * self.eps
                        * torch.eye(bias.size(0), device=bias.device, dtype=bias.dtype)
                    )
                    bias_hessian = jacobian(bias.grad, bias) + diag_comp

                    state["b_ic"] = torch.linalg.cholesky(bias_hessian)

            if group["layer_type"] in ["Linear", "Conv2d"]:

                # Update convariances and inverses
                if update_stats:
                    if self._iteration_counter % self.update_freq == 0:
                        self._compute_covs(group, state)
                        ixxt, iggt, _, _ = self._inv_covs(
                            state["xxt"], state["ggt"], state["num_locations"]
                        )
                        state["ixxt"] = ixxt
                        state["iggt"] = iggt
                    else:
                        if self.alpha != 1:
                            self._compute_covs(group, state)
                if update_params:
                    # Preconditionning
                    gw, gb = self._precond(weight, bias, group, state)
                    # Updating gradients
                    if self.constraint_norm:
                        fisher_norm += (weight.grad * gw).sum()
                    weight.grad.data = gw
                    if bias is not None:
                        if self.constraint_norm:
                            fisher_norm += (bias.grad * gb).sum()
                        bias.grad.data = gb
                # Cleaning
                if "x" in self.state[group["mod"]]:
                    del self.state[group["mod"]]["x"]
                if "gy" in self.state[group["mod"]]:
                    del self.state[group["mod"]]["gy"]
        # Eventually scale the norm of the gradients
        if update_params and self.constraint_norm:
            f_scale = (1.0 / fisher_norm) ** 0.5
            for group in self.param_groups:
                for param in group["params"]:
                    param.grad.data *= f_scale
        if update_stats:
            self._iteration_counter += 1