diff --git a/source/MainComponent.hpp b/source/MainComponent.hpp
index 9d3d930..567155b 100644
--- a/source/MainComponent.hpp
+++ b/source/MainComponent.hpp
@@ -126,13 +126,12 @@ private:
 
 	const arma::mat trainingAt(size_t index)
 	{
-		if (m_pLayer->context().n_rows != m_stack->trainingBatch().n_rows)
-		{
-			m_pLayer->setBatch(m_stack->trainingBatch());
-		}
-
-		if (!m_pLayer->context().is_empty())
+		if (m_pLayer->context().n_cols > 0)
 		{
+			if (m_pLayer->context().n_rows != m_stack->trainingBatch().n_rows)
+			{
+				m_pLayer->setBatch(m_stack->trainingBatch());
+			}
 			return arma::join_rows(m_stack->trainingBatch().row(index), m_pLayer->context().row(index));
 		}
 		else