If you save a checkpoint from within a with model.replicate(), then it will checkpoint the replicated variables. This is the correct behavior, but it might be convenient to add a warning that this is happening and ask if the user didn't intend to checkpoint from outside of the replicate().
If you save a checkpoint from within a
with model.replicate()
, then it will checkpoint the replicated variables. This is the correct behavior, but it might be convenient to add a warning that this is happening and ask if the user didn't intend to checkpoint from outside of the replicate().