Update: I’ve reproduced the error with the following minimal working example. Again, this works perfectly fine on my local machine (so I can’t see any problem with the code itself); the error only appears in Colab.
# <same imports as in previous reply>

# load MNIST data into training & validation sets, and convert to tensors:
from keras.datasets import mnist
(x_train, y_train), (x_valid, y_valid) = mnist.load_data()
x_train, y_train, x_valid, y_valid = map(torch.from_numpy, (x_train, y_train, x_valid, y_valid))

# linear layer z=Wx+b, with W,b drawn from normal distributions:
class GaussianLinear(nn.Module):
    def __init__(self, size_in, size_out, var_W, var_b):
        super().__init__()
        self.size_in, self.size_out = size_in, size_out
        self.var_W, self.var_b = var_W/size_in, var_b  # n.b., must scale var_W by layer width!

        # normally distributed weights with mean=0 and variance=var_W/N:
        norm_vec = torch.normal(mean=torch.zeros(size_out*size_in), std=math.sqrt(self.var_W))
        self.weights = nn.Parameter(norm_vec.view(size_out, size_in))
        
        # normally distributed biases with mean=0 and variance=var_b:
        self.bias = nn.Parameter(torch.normal(mean=torch.zeros(size_out), std=math.sqrt(var_b)))

    def forward(self, x):
        prod = torch.mm(x, self.weights.t())  # Wx
        return torch.add(prod, self.bias)     # Wx+b


# now construct a simple toy network, mimicking the construction in our actual code:
my_layer1 = GaussianLinear(784, 10, 1, 1)
my_layer2 = nn.Tanh()
layers = [my_layer1, my_layer2]
model = nn.Sequential(*layers)

# try training it on the first 5 elements of the training set:
model(x_train[:5])

The last line throws the same error as above:
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-24-44ae810cbc3d> in <module>()
----> 1 model(x_train[:5])

3 frames
<ipython-input-20-e26371d4d8d9> in forward(self, x)
     14 
     15     def forward(self, x):
---> 16         prod = torch.mm(x, self.weights.t())  # Wx
     17         return torch.add(prod, self.bias)     # Wx+b

RuntimeError: self must be a matrix

If any experienced Colab users could shed light on why things are breaking here, I’d be extremely grateful. Otherwise I’ll have to go back to running things on my poor abused desktop.