diff --git a/ethosu/regor/architecture/ethosu55/ethos_u55_register_cs_generator.cpp b/ethosu/regor/architecture/ethosu55/ethos_u55_register_cs_generator.cpp
index 62380fa17521610c41b27ea685f928fa790a13c7..f9b4b39b0d60c2ce3fcbb28ab6471d50e1348685 100644
--- a/ethosu/regor/architecture/ethosu55/ethos_u55_register_cs_generator.cpp
+++ b/ethosu/regor/architecture/ethosu55/ethos_u55_register_cs_generator.cpp
@@ -1397,14 +1397,14 @@ void EthosU55RCSGenerator::InsertTransposeCommand(const HLCStripe *stripe, Tempo
                     depth = sliceShape.Depth();
                     slices = 1;
                     ifmStep = ofmStep = 0;
-                    assert(from == 1 && to == 2);
+                    assert((from == ifm.shape.Size() - 3) && (to == ifm.shape.Size() - 2));
                 }
                 else if ( ofm.transpose == TransposeType::NHCW )
                 {
                     depth = 1;
                     slices = ifm.shape.Height();
                     ifmStep = ofmStep = ifm.shape.ElementsWC() * elementSize;
-                    assert(from == 2 && to == 3);
+                    assert((from == ifm.shape.Size() - 2) && (to == ifm.shape.Size() - 1));
                 }
                 else if ( ofm.transpose == TransposeType::NCWH )
                 {
@@ -1413,7 +1413,7 @@ void EthosU55RCSGenerator::InsertTransposeCommand(const HLCStripe *stripe, Tempo
                     slices = ifm.shape.Width();
                     ifmStep = ifm.shape.Depth() * elementSize;
                     ofmStep = ifm.shape.Height() * elementSize;
-                    assert(from == 1 && to == 3);
+                    assert((from == ifm.shape.Size() - 3) && (to == ifm.shape.Size() - 1));
                 }
                 else
                 {