Update pytorch_lightning import to reflect api change for new version.

OATML · Apr 1, 2022 · 51b4946 · 51b4946
1 parent 9572243
commit 51b4946
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/npt/loss.py b/npt/loss.py
@@ -2,7 +2,7 @@
 
 import torch
 import torch.nn as nn
-from pytorch_lightning.metrics.functional import auroc as lightning_auroc
+from torchmetrics.functional import auroc as lightning_auroc
 
 from npt.utils.encode_utils import torch_cast_to_dtype