fix output shape mismatch after CheckpointBlock

bmtrain/block_layer.py

@@ -49,9 +49,11 @@ def forward(ctx, placeholder, block : 'CheckpointBlock', preserve_rng_state, len
 
         if not isinstance(outputs, list) and not isinstance(outputs, tuple):
             outputs = [outputs]
+            len_outputs = 0
         else:
             outputs = list(outputs)
-        return tuple([len(outputs)] + outputs + [hidden_state["tensor"] for hidden_state in inspector.hidden_states])
+            len_outputs = len(outputs)
+        return tuple([len_outputs] + outputs + [hidden_state["tensor"] for hidden_state in inspector.hidden_states])
 
     @staticmethod
     def backward(ctx, _, *grads):
@@ -453,7 +455,7 @@ def __call__(self, *args, **kwargs):
             all_inputs.append(val)
         outputs = OpCheckpointBlock.apply(placeholder, self, True, len(args), *all_inputs)
         len_output = outputs[0]
-        return outputs[1:1+len_output] if len_output > 1 else outputs[1]
+        return outputs[1:1+len_output] if len_output > 0 else outputs[1]
 
     def __getattr__(self,name:str):
         if name=="_module":

tests/test_all.py

@@ -3,6 +3,7 @@
 
 
 tq = tqdm([
+    ("different_output_shape", 1),
     ("load_ckpt", 1),
     ("init_parameters", 1),
     ("init_parameters_multi_gpu", 4),

tests/test_different_output_shape.py

@@ -0,0 +1,50 @@
+import torch
+import bmtrain as bmt
+
+class Block0(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x
+
+class Block1(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return (x,)
+
+class Block2(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return (x, x)
+
+class Block10(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return [x, x, x, x, x, x, x, x, x, x]
+
+if __name__ == "__main__":
+    bmt.init_distributed()
+    x = torch.tensor([1,2,3.])
+
+    b = bmt.CheckpointBlock(Block0())
+    y = b(x)
+    assert isinstance(y, torch.Tensor)
+
+    b = bmt.CheckpointBlock(Block1())
+    y = b(x)
+    assert isinstance(y, tuple) and len(y)==1
+
+    b = bmt.CheckpointBlock(Block2())
+    y = b(x)
+    assert isinstance(y, tuple) and len(y)==2
+
+    b = bmt.CheckpointBlock(Block10())
+    y = b(x)
+    assert isinstance(y, tuple) and len(y)==10