Dataset from torch to torch does not write in storage

[mariogmarq]

When creating a Dataset object from a pytorch dataset it is passed as a generator, meaning that LazyIndexable will store the whole contents of it in memory. This pull request makes a small fix so we can take advantage of the indexing capabilities of the PyTorch dataset improving memory usage.

[github-actions]

Coverage Report

File	Stmts	Miss	Cover	Missing
flex/common
utils.py	75	15	80%	135–151, 189
flex/data
dataset.py	155	28	82%	53–59, 201, 203, 215, 223–239, 289–290, 307–308, 315, 323, 328, 335–336, 398, 414–415
dataset_pt_utils.py	29	2	93%	48, 59
fed_data_distribution.py	182	1	99%	397
lazy_indexable.py	79	2	97%	55, 89
pluggable_datasets.py	62	1	98%	27
flex/datasets
federated_datasets.py	45	26	42%	35, 39–71, 75–89, 102
flexible_datasets.py	11	3	73%	29, 46–47
flex/pool
aggregators.py	50	4	92%	55–56, 145–146
primitives_pt.py	63	50	21%	46, 61–63, 104–122, 158–170, 197–206, 235–244
primitives_tf.py	42	7	83%	74–84
TOTAL	1286	139	89%

Tests	Skipped	Failures	Errors	Time
190	4 💤	0 ❌	0 🔥	2m 8s ⏱️