julius/prototorch
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[BUGFIX] Initializers can handle Dataloaders now

Browse Source
This commit is contained in:
Alexander Engelsberger 2021-05-21 16:00:20 +02:00
parent 14508f0600
commit ee30d4da5b
1 changed files with 7 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
prototorch/components/initializers.py
View File

@ -7,17 +7,18 @@ import torch
from torch.utils.data import DataLoader, Dataset
def parse_data_arg(data):
if isinstance(data, Dataset):
data, labels = next(iter(DataLoader(data, batch_size=len(data))))
elif isinstance(data, DataLoader):
def parse_data_arg(data_arg):
if isinstance(data_arg, Dataset):
data_arg = DataLoader(data_arg, batch_size=len(data_arg))
if isinstance(data_arg, DataLoader):
data = torch.tensor([])
labels = torch.tensor([])
for x, y in data:
for x, y in data_arg:
data = torch.cat([data, x])
labels = torch.cat([labels, y])
else:
data, labels = data
data, labels = data_arg
if not isinstance(data, torch.Tensor):
wmsg = f"Converting data to {torch.Tensor}."
warnings.warn(wmsg)