Hello Arnav!



 arnavs:

I’m concerned that this won’t scale well when we have large numbers of observations per batch, or huge N


The short answer is to pre-compute X_bar.


    def forward(self, x):
        # forward pass
        M = x.shape[0]
        X_bar = torch.mean(x[:, 2:], dim=1, keepdims=True) 
        Y = torch.cat((x[:, 0:2].reshape([M, 2]), X_bar), dim=1) # [1, c, X_bar]
        output = self.Pb(Y, Y)
        return output



output depends on X_bar but not on the individual values of
the components of X.  So if you call .forward() repeatedly (for
example in your training loop) you will be needlessly computing
X_bar multiple times.
Move the computation of X_bar out of .forward() – and thus
out of the loop – to avoid this unnecessary computation.
(In a similar vein, but less important in terms of computation, I
would also move the construction of Y out of .forward().)
Best.
K. Frank