Update test_data.py (#292)

aiqm · Aug 15, 2019 · 5bc9651 · 5bc9651
1 parent 027c21b
commit 5bc9651
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/tests/test_data.py b/tests/test_data.py
@@ -76,8 +76,7 @@ def testTensorShape(self):
             self.assertLessEqual(species.shape[0], batch_size)
             self.assertEqual(len(coordinates.shape), 3)
             self.assertEqual(coordinates.shape[2], 3)
-            self.assertEqual(coordinates.shape[1], coordinates.shape[1])
-            self.assertEqual(coordinates.shape[0], coordinates.shape[0])
+            self.assertEqual(coordinates.shape[:2], species.shape[:2])
             self.assertEqual(len(energies.shape), 1)
             self.assertEqual(coordinates.shape[0], energies.shape[0])