Update

[ghstack-poisoned]
pytorch · Dec 15, 2024 · 8476c5e · 8476c5e
1 parent 2813d4b
commit 8476c5e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sota-implementations/gail/gail.py b/sota-implementations/gail/gail.py
@@ -233,7 +233,7 @@ def update(data, expert_data, num_network_updates=num_network_updates):
 
                 # Update the networks
                 optim.step()
-        return TensorDict(dloss=d_loss, alpha=alpha).detach()
+        return {"dloss": d_loss, "alpha": alpha}
 
     if cfg.compile.compile:
         update = compile_with_warmup(update, warmup=2, mode=compile_mode)