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

[BUGFIX] Stratified functions work on GPU now

Browse Source
This commit is contained in:
Alexander Engelsberger 2021-06-03 13:19:26 +02:00
parent 1e0a8392a2
commit b724a28a6f
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
prototorch/functions/competitions.py
View File

@ -39,7 +39,7 @@ def stratify_with(values: torch.Tensor,
# skip if stratification is trivial
return values
batch_size = values.size()[0]
winning_values = torch.zeros(num_classes, batch_size)
winning_values = torch.zeros(num_classes, batch_size, device=labels.device)
filler = torch.full_like(values.T, fill_value=fill_value)
for i, cl in enumerate(clabels):
matcher = torch.eq(labels.unsqueeze(dim=1), cl)