the batch size in pytorch is always the first dimension so the inputs_features become the dimension after that