diff --git a/src/rbm/entity.py b/src/rbm/entity.py
index a283c40..62e4473 100644
--- a/src/rbm/entity.py
+++ b/src/rbm/entity.py
@@ -6,6 +6,7 @@ from status import Status
 
 class Entity:
 	def __init__(self, shape: tuple[int, int], params: RbmParams):
+		self.shape = shape
 		self.state = RbmState.from_layer_params(shape)
 		self.params = params