Oam Patel

I can do this, but probably in a couple weeks or so.

Hey @danlaudk , I think this would require creating a new test that checks that an equivalent gated MLP in pytorch gets the same result up to torch.allclose.