Merge pull request deepspeedai#2 from EleutherAI/Stella

Stella
determined-ai · Jan 24, 2021 · f42184b · f42184b
2 parents ef8d702 + c33b239
commit f42184b
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/deepspeed/runtime/pipe/module.py b/deepspeed/runtime/pipe/module.py
@@ -570,6 +570,5 @@ def _is_checkpointable(self, funcs):
         if self.__class__.__name__ == 'GPT2ModelPipe':
             return all('ParallelTransformerLayerPipe' in f.__class__.__name__
                        for f in funcs)
-
-        params = [f.parameters() for f in funcs if isinstance(f, torch.nn.Module)]
-        return any(len(list(p)) > 0 for p in params)
+        return all('TransformerBlock' in f.__class__.__name__
+                       for f in funcs)