How to use the nemo.core.Backend.PyTorch function in NEMO

# parser.add_argument('-evalp', '--evalp',
#                     help='evaluation period', default=1)
# parser.add_argument('-an', '--addName',
#                     help='An add name for the save folder', default='')
# parser.add_argument('-eb', '--eval_batch', help='Evaluation Batch_size',
#                     type=int, default=0)

args = parser.parse_args()
DOMAINS = {"attraction": 0, "restaurant": 1, "taxi": 2, "train": 3, "hotel": 4}

if not os.path.exists(args.data_dir):
    raise ValueError(f'Data not found at {args.data_dir}')

work_dir = f'{args.work_dir}/{args.dataset_name.upper()}'

nf = nemo.core.NeuralModuleFactory(backend=nemo.core.Backend.PyTorch,
                                   local_rank=args.local_rank,
                                   optimization_level=args.amp_opt_level,
                                   log_dir=work_dir,
                                   create_tb_writer=True,
                                   files_to_copy=[__file__],
                                   add_time_to_log_dir=True)

data_layer = nemo_nlp.WOZDSTDataLayer(args.data_dir,
                                      DOMAINS,
                                      batch_size=args.batch_size,
                                      mode='train')
src_ids, src_lens, tgt_ids, tgt_lens, gate_labels, turn_domain = data_layer()
vocab_size = len(data_layer._dataset.vocab)
steps_per_epoch = len(data_layer) // args.batch_size

encoder = EncoderRNN(vocab_size,

opt_level)


name = construct_name('Jasper10x5', lr, batch_size, num_gpus, num_epochs,
                      weight_decay)
tb_writer = SummaryWriter(name)

if args.local_rank is not None:
    device = nemo.core.DeviceType.AllGpu
    print('Doing ALL GPU')
else:
    device = nemo.core.DeviceType.GPU

# instantiate Neural Factory with supported backend
neural_factory = nemo.core.NeuralModuleFactory(
    backend=nemo.core.Backend.PyTorch,
    local_rank=args.local_rank,
    optimization_level=opt_level,
    placement=device)

jasper_model_definition = toml.load("../../examples/nemo_asr/jasper10x5.toml")
jasper_model_definition['placement'] = device
labels = jasper_model_definition['labels']['labels']

train_manifest = "/mnt/D1/Data/librispeech/librivox-train-all.json"
#train_manifest = args.train_manifest
val_manifest1 = "/mnt/D1/Data/librispeech/librivox-dev-clean.json"
# val_manifest2 = "/mnt/D1/Data/librispeech/librivox-dev-other.json"
#val_manifest1 = args.val_manifest1


featurizer_config = jasper_model_definition['input']

from .tacotron2_modules import (MakeGate, Tacotron2Loss, Tacotron2Postnet,
                                Tacotron2Decoder, Tacotron2DecoderInfer,
                                Tacotron2Encoder, TextEmbedding)
from .waveglow_modules import WaveGlowNM, WaveGlowInferNM, WaveGlowLoss
from .data_layers import AudioDataLayer
from .parts.helpers import (waveglow_log_to_tb_func,
                            waveglow_process_eval_batch,
                            waveglow_eval_log_to_tb_func,
                            tacotron2_log_to_tb_func,
                            tacotron2_process_eval_batch,
                            tacotron2_process_final_eval,
                            tacotron2_eval_log_to_tb_func)

name = "nemo_tts"
backend = Backend.PyTorch
__version__ = "0.9.0"

type=str)
parser.add_argument("--dataset_name", default='snips-all', type=str)
parser.add_argument("--data_dir",
                    default='data/nlu/snips',
                    type=str)
parser.add_argument("--query", default='please turn on the light', type=str)
parser.add_argument("--work_dir",
                    default='outputs/SNIPS-ALL/20191010-164934/checkpoints',
                    type=str)
parser.add_argument("--amp_opt_level", default="O0",
                    type=str, choices=["O0", "O1", "O2"])
parser.add_argument("--do_lower_case", action='store_false')

args = parser.parse_args()

nf = nemo.core.NeuralModuleFactory(backend=nemo.core.Backend.PyTorch,
                                   optimization_level=args.amp_opt_level,
                                   log_dir=None)

""" Load the pretrained BERT parameters
See the list of pretrained models, call:
nemo_nlp.huggingface.BERT.list_pretrained_models()
"""
pretrained_bert_model = nemo_nlp.huggingface.BERT(
    pretrained_model_name=args.pretrained_bert_model, factory=nf)
tokenizer = BertTokenizer.from_pretrained(args.pretrained_bert_model)
hidden_size = pretrained_bert_model.local_parameters["hidden_size"]

data_desc = JointIntentSlotDataDesc(args.data_dir,
                                    args.do_lower_case,
                                    args.dataset_name)

learning_rate = args.learning_rate
max_steps = args.max_steps
weight_decay = args.weight_decay
momentum = args.momentum
num_gpus = args.num_gpus

if args.tb_folder is None:
    tb_folder = 'resnet50_fp32'
else:
    tb_folder = args.tb_folder

tb_writer = SummaryWriter(tb_folder)

# instantiate Neural Factory with supported backend
neural_factory = nemo.core.NeuralModuleFactory(
    backend=nemo.core.Backend.PyTorch,
    local_rank=args.local_rank,
    # Set this to nemo.core.Optimization.mxprO1
    # if you have Volta or Turing GPU
    optimization_level=nemo.core.Optimization.mxprO0)

resnet = neural_factory.get_module(name="resnet50",
                                   params={"placement": device},
                                   collection="torchvision",
                                   pretrained=False)

dl_train = neural_factory.get_module(
    name="ImageFolderDataLayer", collection="torchvision",
    params={"batch_size": batch_size,
            "input_size": resnet.input_ports["x"].axis2type[2].dim,
            "shuffle": True,
            "path": args.data_root + "train",

'-1' - step checkpoint won't be saved")
parser.add_argument("--save_step_freq", default=200, type=int,
                    help="Frequency of saving checkpoint \
                    '-1' - step checkpoint won't be saved")
parser.add_argument("--loss_step_freq", default=250, type=int,
                    help="Frequency of printing loss")
parser.add_argument("--use_weighted_loss_punct", action='store_true',
                    help="Flag to indicate whether to use weighted loss \
                    to mitigate classs unbalancing for the punctuation task")

args = parser.parse_args()

if not os.path.exists(args.data_dir):
    raise FileNotFoundError("Dataset not found.")

nf = nemo.core.NeuralModuleFactory(backend=nemo.core.Backend.PyTorch,
                                   local_rank=args.local_rank,
                                   optimization_level=args.amp_opt_level,
                                   log_dir=args.work_dir,
                                   create_tb_writer=True,
                                   files_to_copy=[__file__],
                                   add_time_to_log_dir=True)

nf.logger.info(args)

output_file = f'{nf.work_dir}/output.txt'

if args.bert_checkpoint is None:
    """ Use this if you're using a standard BERT model.
    To see the list of pretrained models, call:
    nemo_nlp.huggingface.BERT.list_pretrained_models()
    """

"attn_model": 'dot',
    "hidden_size": 512,
    "encoder_n_layers": 2,
    "decoder_n_layers": 2,
    "dropout": 0.1,
    "voc_size": 6104 + 3,
    "batch_size": 128,
    "num_epochs": 15,
    "optimizer_kind": "adam",
    "learning_rate": 0.0003,
    "tb_log_dir": "ChatBot",
}

# instantiate Neural Factory with supported backend
neural_factory = nemo.core.NeuralModuleFactory(
    backend=nemo.core.Backend.PyTorch,
    local_rank=None)

# instantiate necessary neural modules
dl = neural_factory.get_module(
    name="DialogDataLayer", collection="tutorials",
    params=config)

# Instance one on EncoderRNN
encoder1 = neural_factory.get_module(
    name="EncoderRNN", collection="tutorials",
    params=config)
# Instance two on EncoderRNN. It will have different weights from instance one
encoder2 = neural_factory.get_module(
    name="EncoderRNN", collection="tutorials",
    params=config)
mixer = neural_factory.get_module(

type=str)
parser.add_argument("--data_dir", default='data/sc/aclImdb', type=str)
parser.add_argument("--dataset_name", default='imdb', type=str)
parser.add_argument("--work_dir", default='outputs', type=str)
parser.add_argument("--save_epoch_freq", default=1, type=int)
parser.add_argument("--save_step_freq", default=-1, type=int)
parser.add_argument("--optimizer_kind", default="adam", type=str)
parser.add_argument("--amp_opt_level", default="O0",
                    type=str, choices=["O0", "O1", "O2"])
parser.add_argument("--do_lower_case", action='store_false')
parser.add_argument("--shuffle_data", action='store_false')

args = parser.parse_args()

work_dir = f'{args.work_dir}/{args.dataset_name.upper()}'
nf = nemo.core.NeuralModuleFactory(backend=nemo.core.Backend.PyTorch,
                                   local_rank=args.local_rank,
                                   optimization_level=args.amp_opt_level,
                                   log_dir=work_dir,
                                   create_tb_writer=True,
                                   files_to_copy=[__file__],
                                   add_time_to_log_dir=True)

""" Load the pretrained BERT parameters
See the list of pretrained models, call:
nemo_nlp.huggingface.BERT.list_pretrained_models()
"""
pretrained_bert_model = nemo_nlp.huggingface.BERT(
    pretrained_model_name=args.pretrained_bert_model, factory=nf)
hidden_size = pretrained_bert_model.local_parameters["hidden_size"]
tokenizer = BertTokenizer.from_pretrained(args.pretrained_bert_model)

args = parse_args()

    name = construct_name(
        args.exp_name,
        args.lr,
        args.batch_size,
        args.num_epochs,
        args.weight_decay,
        args.optimizer)
    work_dir = name
    if args.work_dir:
        work_dir = os.path.join(args.work_dir, name)

    # instantiate Neural Factory with supported backend
    neural_factory = nemo.core.NeuralModuleFactory(
        backend=nemo.core.Backend.PyTorch,
        local_rank=args.local_rank,
        optimization_level=args.amp_opt_level,
        log_dir=work_dir,
        checkpoint_dir=args.checkpoint_dir,
        create_tb_writer=args.create_tb_writer,
        files_to_copy=[args.model_config, __file__],
        cudnn_benchmark=args.cudnn_benchmark,
        tensorboard_dir=args.tensorboard_dir)
    args.num_gpus = neural_factory.world_size

    logger = neural_factory.logger
    args.checkpoint_dir = neural_factory.checkpoint_dir

    if args.local_rank is not None:
        logger.info('Doing ALL GPU')

name = construct_name(
        args.exp_name,
        args.lr,
        args.batch_size,
        args.max_steps,
        args.num_epochs,
        args.weight_decay,
        args.optimizer,
        args.iter_per_step)
    log_dir = name
    if args.work_dir:
        log_dir = os.path.join(args.work_dir, name)

    # instantiate Neural Factory with supported backend
    neural_factory = nemo.core.NeuralModuleFactory(
        backend=nemo.core.Backend.PyTorch,
        local_rank=args.local_rank,
        optimization_level=args.amp_opt_level,
        log_dir=log_dir,
        checkpoint_dir=args.checkpoint_dir,
        create_tb_writer=args.create_tb_writer,
        files_to_copy=[args.model_config, __file__],
        cudnn_benchmark=args.cudnn_benchmark,
        tensorboard_dir=args.tensorboard_dir)
    args.num_gpus = neural_factory.world_size

    logger = neural_factory.logger
    checkpoint_dir = neural_factory.checkpoint_dir
    if args.local_rank is not None:
        logger.info('Doing ALL GPU')

    # build dags