diff --git a/src/test_utils/models.jl b/src/test_utils/models.jl
index 4aa2aaa42..d11158d1a 100644
--- a/src/test_utils/models.jl
+++ b/src/test_utils/models.jl
@@ -406,12 +406,14 @@ end
     s .~ InverseGamma(2, 3)
     m = TV(undef, 2)
     m .~ Normal.(0, sqrt.(s))
-    return s, m
+    return (; s=s, m=m)
 end
 
 @model function demo_assume_submodel_observe_index_literal()
     # Submodel prior
-    @submodel s, m = _prior_dot_assume()
+    vars ~ to_submodel(_prior_dot_assume())
+    s = vars.s
+    m = vars.m
     1.5 ~ Normal(m[1], sqrt(s[1]))
     2.0 ~ Normal(m[2], sqrt(s[2]))
 
@@ -449,7 +451,7 @@ end
     m .~ Normal.(0, sqrt.(s))
 
     # Submodel likelihood
-    @submodel _likelihood_mltivariate_observe(s, m, x)
+    _ignore ~ to_submodel(_likelihood_mltivariate_observe(s, m, x))
 
     return (; s=s, m=m, x=x, logp=getlogp(__varinfo__))
 end