Method compile

deepspeed/runtime/pipe/module.py:692–698

(self, *args, **kwargs)

        return None

    def compile(self, *args, **kwargs):
        for idx, layer in enumerate(self.forward_funcs):
            if isinstance(layer, nn.Module):
                layer.compile(*args, **kwargs)
            else:
                new_layer = torch.compile(layer, *args, **kwargs)
                self.forward_funcs[idx] = new_layer