Fix issue where matrix losing dimension