move an assert

pi_zero_pytorch/pi_zero.py

@@ -664,7 +664,6 @@ def forward_with_reward_cfg(
         keep_parallel_frac = 0.,
         **kwargs
     ):
-        assert self.can_cfg, 'you need to train with reward token dropout'
 
         with_reward_cache, without_reward_cache = cached_state_keys_values
 
@@ -686,6 +685,8 @@ def forward_with_reward_cfg(
         if not exists(reward_tokens) or cond_scale == 0.:
             return action_flow_with_reward, (with_reward_cache_kv, None)
 
+        assert self.can_cfg, 'you need to train with reward token dropout'
+
         action_flow_without_reward, without_reward_cache_kv = self.forward(
             *args,
             cached_state_keys_values = without_reward_cache,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pi-zero-pytorch"
-version = "0.0.19"
+version = "0.0.20"
 description = "π0 in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }