Repository: continualai/avalanche
Issue: bug AvalancheDataset with multi-task labels

```
import torch
from avalanche.benchmarks.utils import AvalancheTensorDataset
x = torch.rand(32, 10)
y = torch.rand(32, 10)
t = torch.ones(32)  # single task
data = AvalancheTensorDataset(x, y, targets=1, task_labels=t)
print(len(data.task_set))
```
expected: 1
result: 32

Changed files:
avalanche/benchmarks/utils/avalanche_dataset.py, tests/test_custom_streams.py, tests/test_avalanche_dataset.py

diff --git a/avalanche/benchmarks/utils/avalanche_dataset.py b/avalanche/benchmarks/utils/avalanche_dataset.py
index 72caa55c..0236b3a1 100644
--- a/avalanche/benchmarks/utils/avalanche_dataset.py
+++ b/avalanche/benchmarks/utils/avalanche_dataset.py
@@ -835,7 +835,8 @@ class AvalancheDataset(IDatasetWithTargets[T_co, TTargetType], Dataset[T_co]):
                     'Invalid amount of task labels. It must be equal to the '
                     'number of patterns in the dataset. Got {}, expected '
                     '{}!'.format(len(task_labels), len(dataset)))
-            return task_labels
+
+            return SubSequence(task_labels, converter=int)
 
         return _make_task_labels_from_supported_dataset(dataset)
 
@@ -1406,7 +1407,7 @@ class AvalancheTensorDataset(AvalancheDataset[T_co, TTargetType]):
                                                    YTransform]] = None,
                  initial_transform_group: str = 'train',
                  task_labels: Union[int, Sequence[int]] = None,
-                 targets: Sequence[TTargetType] = None,
+                 targets: Union[Sequence[TTargetType], int] = None,
                  dataset_type: AvalancheDatasetType =
                  AvalancheDatasetType.UNDEFINED,
                  collate_fn: Callable[[List], Any] = None,
@@ -1730,7 +1731,7 @@ class AvalancheConcatDataset(AvalancheDataset[T_co, TTargetType]):
                     'Invalid amount of task labels. It must be equal to the '
                     'number of patterns in the dataset. Got {}, expected '
                     '{}!'.format(len(task_labels), self._overall_length))
-            return task_labels
+            return SubSequence(task_labels, converter=int)
 
         concat_t_labels = []
         for dataset_idx, single_dataset in enumerate(self._dataset_list):
diff --git a/tests/test_avalanche_dataset.py b/tests/test_avalanche_dataset.py
index ac543994..32a41773 100644
--- a/tests/test_avalanche_dataset.py
+++ b/tests/test_avalanche_dataset.py
@@ -218,6 +218,44 @@ class AvalancheDatasetTests(unittest.TestCase):
         with self.assertRaises(KeyError):
             subset_task0 = dataset.task_set[0]
 
+    def test_avalanche_dataset_tensor_task_labels(self):
+        x = torch.rand(32, 10)
+        y = torch.rand(32, 10)
+        t = torch.ones(32)  # Single task
+        dataset = AvalancheTensorDataset(x, y, targets=1, task_labels=t)
+
+        x2, y2, t2 = dataset[:]
+
+        self.assertIsInstance(x2, Tensor)
+        self.assertIsInstance(y2, Tensor)
+        self.assertIsInstance(t2, Tensor)
+        self.assertTrue(torch.equal(x, x2))
+        self.assertTrue(torch.equal(y, y2))
+        self.assertTrue(torch.equal(t.to(int), t2))
+
+        self.assertListEqual([1] * 32,
+                             list(dataset.targets_task_labels))
+
+        # Regression test for #654
+        self.assertEqual(1, len(dataset.task_set))
+
+        subset_task1 = dataset.task_set[1]
+        self.assertIsInstance(subset_task1, AvalancheDataset)
+        self.assertEqual(len(dataset), len(subset_task1))
+
+        with self.assertRaises(KeyError):
+            subset_task0 = dataset.task_set[0]
+
+        with self.assertRaises(KeyError):
+            subset_task0 = dataset.task_set[2]
+
+        # Check single instance types
+        x2, y2, t2 = dataset[0]
+
+        self.assertIsInstance(x2, Tensor)
+        self.assertIsInstance(y2, Tensor)
+        self.assertIsInstance(t2, int)
+
     def test_avalanche_dataset_uniform_task_labels_simple_def(self):
         dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                               download=True)
diff --git a/tests/test_custom_streams.py b/tests/test_custom_streams.py
index 51151d7f..a75b1a85 100644
--- a/tests/test_custom_streams.py
+++ b/tests/test_custom_streams.py
@@ -80,7 +80,7 @@ class CustomStreamsTests(unittest.TestCase):
 
             self.assertTrue(torch.equal(expect_x, got_x))
             self.assertTrue(torch.equal(expect_y, got_y))
-            self.assertTrue(torch.equal(expect_t, got_t))
+            self.assertEqual(int(expect_t), got_t)
 
             exp_t_labels = set(exp.task_labels)
             self.assertLess(max(exp_t_labels), 5)
@@ -92,7 +92,7 @@ class CustomStreamsTests(unittest.TestCase):
 
             self.assertTrue(torch.equal(expect_x, got_x))
             self.assertTrue(torch.equal(expect_y, got_y))
-            self.assertTrue(torch.equal(expect_t, got_t))
+            self.assertEqual(int(expect_t), got_t)
 
             exp_t_labels = set(exp.task_labels)
             self.assertLess(max(exp_t_labels), 3)
@@ -104,7 +104,7 @@ class CustomStreamsTests(unittest.TestCase):
 
             self.assertTrue(torch.equal(expect_x, got_x))
             self.assertTrue(torch.equal(expect_y, got_y))
-            self.assertTrue(torch.equal(expect_t, got_t))
+            self.assertEqual(int(expect_t), got_t)
 
             exp_t_labels = set(exp.task_labels)