diff --git a/source/MainComponent.hpp b/source/MainComponent.hpp
index 3f2be79..9d3d930 100644
--- a/source/MainComponent.hpp
+++ b/source/MainComponent.hpp
@@ -126,7 +126,7 @@ private:
 
 	const arma::mat trainingAt(size_t index)
 	{
-		if (m_pLayer->context().is_empty())
+		if (m_pLayer->context().n_rows != m_stack->trainingBatch().n_rows)
 		{
 			m_pLayer->setBatch(m_stack->trainingBatch());
 		}