Fixes flatten on dim=1

[mmcdermott]

Closes #39

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 90.78%. Comparing base (eaaa6bb) to head (5cc8569).
Report is 3 commits behind head on dev.

✅ All tests successful. No failed tests found.

Additional details and impacted files

@@            Coverage Diff             @@
##              dev      #40      +/-   ##
==========================================
+ Coverage   90.74%   90.78%   +0.04%     
==========================================
  Files           2        2              
  Lines         443      445       +2     
==========================================
+ Hits          402      404       +2     
  Misses         41       41              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.