Hi,
The obvious thing you can do to speed things up is to remove the two for loops and make all operations work directly on tensors:
for i in range(seq_len):
        for node in range(numNodes):
            mean = torch.cat((o_mux[i, node], o_muy[i, node]))
            sigma_xx = o_sx[i, node]*o_sx[i, node]
            sigma_yy = o_sy[i, node]*o_sy[i, node] 
            sigma_xy = o_corr[i, node]*o_sx[i, node]*o_sy[i, node]
            sigma_yx = o_corr[i, node]*o_sy[i, node]*o_sx[i, node]

Can just become
mean = torch.cat((o_mux, o_muy))
sigma_xx = o_sx *o_sx
sigma_yy = o_sy * o_sy 
sigma_xy = o_corr * o_sx * o_sy
sigma_yx = o_corr * o_sy * o_sx