diff --git a/src/tests/test_learn_encoded_labels.py b/src/tests/test_learn_encoded_labels.py
index aa57bda..8aefde8 100644
--- a/src/tests/test_learn_encoded_labels.py
+++ b/src/tests/test_learn_encoded_labels.py
@@ -9,7 +9,7 @@ from rbm.matrix import Mat, np
 class LabelLearner(Model):
 	def __init__(self, name: str, dim, work_dir: str = '.'):
 		super().__init__(name, work_dir)
-		self.unit1 = Entity(dim, EntityParams(do_gaussian_hidden=False), TrainingParams(learning_rate=0.01, momentum=0.9, do_rao_blackwell=True, num_epochs=10000))
+		self.unit1 = Entity(dim, EntityParams(do_gaussian_hidden=False), TrainingParams(learning_rate=0.01, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))
 
 	def forward(self, x: Mat):
 		x = self.unit1.forward(x)
@@ -66,13 +66,13 @@ if __name__ == "__main__":
 		recon = model.backward(hidden)
 		img_hidden = np.reshape(hidden, (4, 4))
 		img_recon = np.reshape(recon, (label_w, label_h))
-		axes[0][index].imshow(img_hidden, cmap=cmap)
+		axes[0][index].imshow(img_hidden.get(), cmap=cmap)
 		axes[0][index].axis('off')
 		axes[0][index].set_title(f'{test_labels[index]}')
-		axes[1][index].imshow(img_recon, cmap=cmap)
+		axes[1][index].imshow(img_recon.get(), cmap=cmap)
 		axes[1][index].axis('off')
 		axes[1][index].set_title(f'{test_labels[index]}')
-		axes[2][index].imshow(np.reshape(inp, (label_w, label_h)), cmap=cmap)
+		axes[2][index].imshow(np.reshape(inp.get(), (label_w, label_h)), cmap=cmap)
 		axes[2][index].axis('off')
 		axes[2][index].set_title(f'{test_labels[index]}')
 
diff --git a/src/tests/test_linear.py b/src/tests/test_linear.py
index abfafde..1796eb7 100644
--- a/src/tests/test_linear.py
+++ b/src/tests/test_linear.py
@@ -40,7 +40,7 @@ def linear():
 	model.init(0.1)
 
 	# Load weights (if exists)
-#	model.load()
+	model.load()
 
 	# Prepare training data
 	training_batch = np.random.randn(N_CASES, N_VIS, dtype=np.float64)