I was a bit confused whether nn.Dense support multi-dimensional inputs through flatten=False?
The doc says:
Note: the input must be a tensor with rank 2. Use flatten to convert it to rank 2 manually if necessary.
flatten (bool) – Whether the input tensor should be flattened. If true, all but the first axis of input data are collapsed together. If false, all but the last axis of input data are kept the same, and the transformation applies on the last axis.
In my experience:
I always use dense with flatten just to eliminate potential issues. I have never tested the flatten=False option to my satisfaction, but I may just lack trying.