use wandb.watch

wandb · Aug 16, 2022 · 15ede5d · 15ede5d
1 parent c9c6abd
commit 15ede5d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/unit_tests/test_torch_full.py b/tests/unit_tests/test_torch_full.py
@@ -160,7 +160,7 @@ def test_all_logging(relay_server, wandb_init):
     with relay_server() as relay:
         run = wandb_init()
         net = ConvNet()
-        run.watch(
+        wandb.watch(
             net,
             log="all",
             log_freq=1,
@@ -188,7 +188,7 @@ def test_embedding_dict_watch(relay_server, wandb_init):
     with relay_server() as relay:
         run = wandb_init()
         model = EmbModelWrapper()
-        run.watch(model, log_freq=1, idx=0)
+        wandb.watch(model, log_freq=1, idx=0)
         opt = torch.optim.Adam(params=model.parameters())
         inp = torch.randint(16, [8, 5])
         out = model(inp)