Fix device allocation

[Kuzaiherba]

Description

This PR fixes device allocation issues in PyTorch tensor creation within the metrics module. The fix ensures that newly created weight tensors are placed on the same device as the target tensors, preventing potential device mismatch errors during tensor operations.

Example / Current workflow

Include a sample workflow to either (a) reproduce the bug with current codebase or (b) showcase the deficiency does this PR seeks to address

from chemprop.nn import BinaryClassificationFFN
from chemprop.nn.metrics import BCELoss


task_dim = 3
predictor = BinaryClassificationFFN(
    n_tasks=task_dim,
    criterion=BCELoss(task_weights=torch.ones(task_dim)),
)

device = "mps"
predictor.to(device=device)
predictor.criterion(
    torch.ones((10, task_dim)).to(device=device),
    torch.ones((10, task_dim)).to(device=device)
)
>>> RuntimeError: Encountered different devices in metric calculation (see stacktrace for details). This could be due to the metric class not being on the same device as input. Instead of `metric=BCELoss(...)` try to do `metric=BCELoss(...).to(device)` where device corresponds to the device of the input.

Bugfix / Desired workflow

Include either (a) the same workflow from above with the correct output produced via this PR (b) some (pseudo)code containing the new workflow that this PR will (seek to) implement

from chemprop.nn import BinaryClassificationFFN
from chemprop.nn.metrics import BCELoss


task_dim = 3
predictor = BinaryClassificationFFN(
    n_tasks=task_dim,
    criterion=BCELoss(task_weights=torch.ones(task_dim)),
)

device = "mps"
predictor.to(device=device)
predictor.criterion(
    torch.ones((10, task_dim)).to(device=device),
    torch.ones((10, task_dim)).to(device=device)
)
>>> Out[1]: tensor(0.3133, device='mps:0')

Questions

If there are open questions about implementation strategy or scope of the PR, include them here

Relevant issues

If appropriate, please tag them here and include a quick summary

Checklist

• linted with flake8?
• (if appropriate) unit tests added?

[Copilot]

Pull Request Overview

This PR fixes device allocation issues in PyTorch tensor creation within the metrics module. The fix ensures that newly created weight tensors are placed on the same device as the target tensors, preventing potential device mismatch errors during tensor operations.

• Added explicit device specification to torch.ones() calls for weight tensor creation
• Applied the fix consistently across three different metric update methods
• Maintained the same conditional logic while ensuring device compatibility

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}