Scalar inputs to PyTorch model

[ibengtsson]

Hi! Do you have any advice on how to best supply scalar inputs to a model? I want to do a mean-pooling operation (averaging over the rows in a matrix), but the matrix will be zero-padded to a fixed size and I would like to supply the number of non-zero rows as an input.

For context, I have a matrix x of size (n, m) and want to take the average over the k non-zero elements in each row. My idea was to use a matrix multiplication (which I've already implemented for another operation, using the extension API), multiplying a binary vector y of size (1, n) with x to do the summation. This would give y * x = z, with z having a size (1, m) as expected. However, to actually get an average and not just a sum I need to be able to divide each element of z with k...

Preferably I'd like to use PyTorch. I'd be grateful for any advice on how to best proceed!