diff --git a/nvpl_tensor/contraction/contraction.c b/nvpl_tensor/contraction/contraction.c
index f58ee1c..a67ec90 100644
--- a/nvpl_tensor/contraction/contraction.c
+++ b/nvpl_tensor/contraction/contraction.c
@@ -67,17 +67,17 @@ int main()
     int64_t elementsA = 1;
     for (int i = 0; i < nmodeA; ++i)
     {
-        elementsA *= extent[i];
+        elementsA *= extentA[i];
     }
     int64_t elementsB = 1;
     for (int i = 0; i < nmodeB; ++i)
     {
-        elementsB *= extent[i];
+        elementsB *= extentB[i];
     }
     int64_t elementsC = 1;
     for (int i = 0; i < nmodeC; ++i)
     {
-        elementsC *= extent[i];
+        elementsC *= extentC[i];
     }
 
     int64_t sizeA = sizeof(floatTypeA) * elementsA;