Fix incorrectly written assert. This assert would pass if the
work_dim was 1 but would fail for larger dimensions. We should add
a test case!