There is no nn.Module in torch to perform tensor slicing, e.g. to implement
certain ResNets for CIFAR10. This PR adds a custom slicing module with support for some second-order extensions.
Add Slicing module
Implement jac_t_mat_prod
Support {Batch}DiagGGN{Exact, MC}, {Batch}DiagHessian, SqrtGGN{Exact, MC} extensions
There is no
nn.Module
intorch
to perform tensor slicing, e.g. to implement certain ResNets for CIFAR10. This PR adds a custom slicing module with support for some second-order extensions.Slicing
modulejac_t_mat_prod
{Batch}DiagGGN{Exact, MC}
,{Batch}DiagHessian
,SqrtGGN{Exact, MC}
extensions