Skip to content

Commit

Permalink
Move targets tensor to GPU for eval
Browse files Browse the repository at this point in the history
  • Loading branch information
SecroLoL committed Jan 12, 2024
1 parent a3ec727 commit 023dd95
Showing 1 changed file with 1 addition and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ def train(self, num_epochs: int, save_name: str, args: Mapping, eval_file: str,
if self.weighted_loss: # BCEWithLogitsLoss requires a vector for target where probability is 1 on the true label class, and 0 on others.
targets = torch.stack([torch.tensor([1, 0]) if label == 0 else torch.tensor([0, 1]) for label in labels]).to(dtype=torch.float32).to(device)
else: # CELoss accepts target as just raw label
targets = labels
targets = labels.to(device)

loss = self.criterion(outputs, targets)

Expand Down

0 comments on commit 023dd95

Please sign in to comment.