[tensor] fix a assertion in colo_tensor cross_entropy (#1232)

colossalai/nn/_ops/loss.py

@@ -23,7 +23,7 @@ def colo_cross_entropy(input_tensor: GeneralTensor,
     input_tensor = convert_to_colo_tensor(input_tensor, pg)
 
     if input_tensor.is_replicate():    # Input is gathered
-        assert target.is_replicate() and weight.is_replicate(), \
+        assert target.is_replicate() and (weight is None or weight.is_replicate()), \
             "Target tensor and weight tensor both should be complete"
         output = F.cross_entropy(input_tensor,
                                  target,

tests/test_tensor/test_loss_func.py

@@ -41,7 +41,7 @@ def run_dist(rank, world_size, port):
 
 
 @pytest.mark.dist
-@pytest.mark.parametrize('world_size', [2])
+@pytest.mark.parametrize('world_size', [1, 2])
 @rerun_if_address_is_in_use()
 def test_loss_func(world_size):
     run_func = partial(run_dist, world_size=world_size, port=free_port())
@@ -49,4 +49,4 @@ def test_loss_func(world_size):
 
 
 if __name__ == '__main__':
-    test_loss_func(2)
+    test_loss_func(1)