diff --git a/homeworks/assignment05_bagging_and_oob/assignment_bagging_and_oob.ipynb b/homeworks/assignment05_bagging_and_oob/assignment_bagging_and_oob.ipynb
index e20e493a2..1a3119a68 100644
--- a/homeworks/assignment05_bagging_and_oob/assignment_bagging_and_oob.ipynb
+++ b/homeworks/assignment05_bagging_and_oob/assignment_bagging_and_oob.ipynb
@@ -69,8 +69,7 @@
     "        self.data = None\n",
     "        self.target = None\n",
     "        self._generate_splits(data)\n",
-    "        assert len(set(list(map(len, self.indices_list)))) == 1, 'All bags should be of the same length!'\n",
-    "        assert list(map(len, self.indices_list))[0] == len(data), 'All bags should contain `len(data)` number of elements!'\n",
+    "        assert all(len(indices) == len(data) for indices in self.indices_list), 'All bags should contain `len(data)` number of elements!'\n",
     "        self.models_list = []\n",
     "        for bag in range(self.num_bags):\n",
     "            model = model_constructor()\n",
diff --git a/homeworks/assignment05_bagging_and_oob/bagging.py b/homeworks/assignment05_bagging_and_oob/bagging.py
index a7952f9dd..71e599547 100644
--- a/homeworks/assignment05_bagging_and_oob/bagging.py
+++ b/homeworks/assignment05_bagging_and_oob/bagging.py
@@ -27,8 +27,7 @@ def fit(self, model_constructor, data, target):
         self.data = None
         self.target = None
         self._generate_splits(data)
-        assert len(set(list(map(len, self.indices_list)))) == 1, 'All bags should be of the same length!'
-        assert list(map(len, self.indices_list))[0] == len(data), 'All bags should contain `len(data)` number of elements!'
+        assert all(len(indices) == len(data) for indices in self.indices_list), 'All bags should contain `len(data)` number of elements!'
         self.models_list = []
         for bag in range(self.num_bags):
             model = model_constructor()