diff --git a/source/MainComponent.cpp b/source/MainComponent.cpp
index 13fca28..e3534db 100644
--- a/source/MainComponent.cpp
+++ b/source/MainComponent.cpp
@@ -946,3 +946,4 @@ void MainComponent::updateControls()
 	patterSlider->setValue(m_trainingIndex);
 
 }
+
diff --git a/source/MainComponent.hpp b/source/MainComponent.hpp
index e88e4dc..4445dfc 100644
--- a/source/MainComponent.hpp
+++ b/source/MainComponent.hpp
@@ -130,7 +130,15 @@ private:
 		{
 			m_pLayer->setBatch(m_stack->trainingData());
 		}
-		return arma::join_rows(m_stack->trainingData().row(index), m_pLayer->context().row(index));
+
+		if (!m_pLayer->context().is_empty())
+		{
+			return arma::join_rows(m_stack->trainingData().row(index), m_pLayer->context().row(index));
+		}
+		else
+		{	
+			return m_stack->trainingData().row(index);
+		}
 	}
 	
 	void updateControls();