ahrefs
diff --git a/‎lib/train.ml‎
Lines changed: 4 additions & 3 deletions b/‎lib/train.ml‎
Lines changed: 4 additions & 3 deletions
@@ -488,21 +488,22 @@ let example_train_loop ?(disable_rootness_check = false) ~seed ~batch_size ~init
   }
 
 (* Note: this will get nicer with modular explicits. *)
-let%track3_sexp forward_and_ctx ?(disable_rootness_check = false)
+let%track3_sexp forward_and_ctx ?(hosted=true) ?(disable_rootness_check = false)
     (type buffer_ptr dev runner event optimize_ctx)
     (module Backend : Backend
       with type buffer_ptr = buffer_ptr
        and type dev = dev
        and type runner = runner
        and type optimize_ctx = optimize_ctx
        and type event = event) ctx ?(bindings = IDX.empty) t =
+  if hosted then set_hosted t.Tensor.value;
   let routine =
     Backend.(link ctx @@ compile ctx.optimize_ctx bindings @@ forward ~disable_rootness_check t)
   in
   if not disable_rootness_check then Tensor.remove_bprop_root t;
   Task.run routine.schedule;
   routine.context
 
-let forward_and_forget ?disable_rootness_check backend ctx ?bindings t =
+let forward_and_forget ?hosted ?disable_rootness_check backend ctx ?bindings t =
   (* FIXME: to properly forget we need to free the incrementally-allocated memory! *)
-  ignore @@ forward_and_ctx ?disable_rootness_check backend ctx ?bindings t
+  ignore @@ forward_and_ctx ?hosted ?disable_rootness_check backend ctx ?bindings t