Fixed another bug

src/main/resources/templates/gluon/saveAttentionImageTest.ftl

@@ -11,7 +11,7 @@
                             attention_resized = np.resize(attention.asnumpy(), (8, 8))
                             ax = fig.add_subplot(max_length//3, max_length//4, l+1)
                             ax.set_title(dict[int(mx.nd.slice_axis(mx.nd.argmax(outputs[l+1], axis=1), axis=0, begin=0, end=1).asscalar())])
-                            img = ax.imshow(test_images[0+batch_size*(batch_i)])
+                            img = ax.imshow(test_images[0+test_batch_size*(batch_i)])
                             ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())
 
 

src/main/resources/templates/gluon/saveAttentionImageTrain.ftl

@@ -17,7 +17,7 @@
                             attention_resized = np.resize(attention.asnumpy(), (8, 8))
                             ax = fig.add_subplot(max_length//3, max_length//4, l+1)
                             ax.set_title(dict[int(labels[l+1][0].asscalar())])
-                            img = ax.imshow(train_images[0+batch_size*(batch_i)])
+                            img = ax.imshow(train_images[0+test_batch_size*(batch_i)])
                             ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())