Remove outdated TP restriction for zero_outside_topk

Signed-off-by: Zhaopeng Qiu <qiuzhaopeng@foxmail.com>

Author: zpqiu

nemo_rl/algorithms/distillation.py

@@ -153,13 +153,7 @@ def setup(
     assert generation_config is not None, (
         "A generation config in the PolicyConfig is required for distillation"
     )
-    assert (  # [TODO] we may support this for tp in the future
-        not loss_config.get("zero_outside_topk", False)
-        or (policy_config["dtensor_cfg"]["tensor_parallel_size"] == 1)
-    ), (
-        f"zero_outside_topk=True requires tensor_parallel_size=1, "
-        f"but got tensor_parallel_size={policy_config['dtensor_cfg']['tensor_parallel_size']}. "
-    )
+
     # Set random seed
     set_seed(distillation_config["seed"])