Make differentiating inputs optional

[Xuzzo]

Hello,

I am calculating the EK-FAC representation on the linear layers of a model with an embedding layer. The inputs are integers which are lookup values and therefore should not be differentiated through. Within the calculation of the kfac blocks in the Jacobian class the inputs are labelled as requiring grad, which gives me RuntimeError: only Tensors of floating point and complex dtype can require gradients. As far as I can tell, this is just so that the hooks can be triggered on all the layers, so a simple backward call would be enough.

Here is a simple code to replicate the error

from torch.nn import Embedding, Linear

from nngeometry.metrics import FIM
from nngeometry.object import PMatEKFAC
import torch as th
from nngeometry.layercollection import LayerCollection


class EmbeddingModel(th.nn.Module):

    def __init__(
        self,
        n_input: int,
        n_output: int,
    ):
        super().__init__()
        self.embedding = Embedding(n_input, 10)
        self.fc1 = Linear(10, n_output)

    def forward(self, x):
        return self.fc1(self.embedding(x))

if __name__ == "__main__":
    model = EmbeddingModel(10, 1)
    active_layers = LayerCollection()
    active_layers.add_layer_from_model(model, model.fc1)
    dataset = th.utils.data.TensorDataset(th.randint(1, 1000, (100,)), th.randint(0, 1, (100,)))
    loader = th.utils.data.DataLoader(dataset, batch_size=10)

    F_ekfac = FIM(model, loader, PMatEKFAC, 1, variant='classif_logits', layer_collection=active_layers)

Would it be possible to remove the inputs.requires_grad = True or make it optional? Thanks a lot for your help and for your very instructive library.