Fix module name in Gaussian RSSM (#52)

danijar · Apr 27, 2023 · c64680b · c64680b
1 parent dab3648
commit c64680b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dreamerv3/nets.py b/dreamerv3/nets.py
@@ -78,7 +78,7 @@ def get_dist(self, state, argmax=False):
     else:
       mean = state['mean'].astype(f32)
       std = state['std'].astype(f32)
-      return tfp.MultivariateNormalDiag(mean, std)
+      return tfd.MultivariateNormalDiag(mean, std)
 
   def obs_step(self, prev_state, prev_action, embed, is_first):
     is_first = cast(is_first)