diff --git a/scripts/py/templates/TestData.cc.template b/scripts/py/templates/TestData.cc.template
index 1acd14d..d0f2698 100644
--- a/scripts/py/templates/TestData.cc.template
+++ b/scripts/py/templates/TestData.cc.template
@@ -32,7 +32,7 @@
 
 const {{data_type}}* get_ifm_data_array(const uint32_t idx)
 {
-    if (idx < NUMBER_OF_FM_FILES) {
+    if (idx < NUMBER_OF_IFM_FILES) {
         return ifm_arrays[idx];
     }
     return nullptr;
@@ -40,7 +40,7 @@
 
 const {{data_type}}* get_ofm_data_array(const uint32_t idx)
 {
-    if (idx < NUMBER_OF_FM_FILES) {
+    if (idx < NUMBER_OF_OFM_FILES) {
         return ofm_arrays[idx];
     }
     return nullptr;