Update Iris dataset

2021-05-17 16:57:13 +02:00 · 2021-05-17 16:57:13 +02:00 · 639198e774
commit 639198e774
parent 768d969f89
1 changed files with 3 additions and 1 deletions
--- a/prototorch/datasets/iris.py
+++ b/prototorch/datasets/iris.py
@ -10,6 +10,8 @@ from sklearn.datasets import load_iris


 class Iris(NumpyDataset):
-    def __init__(self):
+    def __init__(self, dims=None):
        x, y = load_iris(return_X_y=True)
+        if dims:
+            x = x[:, dims]
        super().__init__(x, y)