def inner_update_step(self, fast_params, x, y, t):
        """Update fast weights using current samples and
        return the updated fast model.
        """
        logits = torch.func.functional_call(self.model, fast_params, (x, t))
        loss = self._criterion(logits, y)

        # Compute gradient with respect to the current fast weights
        grads = list(
            torch.autograd.grad(
                loss,
                fast_params.values(),
                retain_graph=self.second_order,
                create_graph=self.second_order,
                allow_unused=True,
            )
        )

        # Clip grad norms
        grads = [
            (
                torch.clamp(g, min=-self.grad_clip_norm, max=self.grad_clip_norm)
                if g is not None
                else g
            )
            for g in grads
        ]

        # New fast parameters
        new_fast_params = {
            n: param - alpha * grad if grad is not None else param
            for ((n, param), alpha, grad) in zip(
                fast_params.items(), self.alpha_params.parameters(), grads
            )
        }

        return new_fast_params
Update fast weights using current samples and
        return the updated fast model.
        
avalanche/training/supervised/lamaml_v2.py