Commit 6106127c authored by Raul Puri's avatar Raul Puri
Browse files

optionally allow no cpu

parent 569b3dab

megatron/initialize.py

+4 −3
Original line number Diff line number Diff line
@@ -29,9 +29,10 @@ from megatron.global_vars import set_global_variables 




def initialize_megatron(extra_args_provider=None, args_defaults={},

                        ignore_unknown_args=False):

                        ignore_unknown_args=False, allow_no_cuda=False):

    """Set global variables, initialize distributed, and

    set autoresume and random seeds."""

    if not allow_no_cuda:

        # Make sure cuda is available.

        assert torch.cuda.is_available(), 'Megatron requires CUDA.'