Error when continuing training from checkpoint
Hello. I'm having trouble when trying to resume training when doing it with FastPitch. I'm using FastPitch_Finetuning.ipynb in the last version (r1.8.2), with a custom dataset. The generated checkpoint is ok and I can do inference with it, but I cannot use it from resuming training. Here is the error:

```
Restored all states from the checkpoint file at /content/drive/MyDrive/TTS/CHECKPOINTS/checkpoints_dataset_mix_92-1_7hours_2_plus15_16/FastPitch/checkpoints/FastPitch--v_loss=1.8955-epoch=274-last.ckpt
Epoch 274:   0% 0/351 [00:00<00:00, -4587799.47it/s][NeMo W 2022-05-09 13:39:54 nemo_logging:349] /usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py:136: UserWarning: You're resuming from a checkpoint that ended before the epoch ended. This can cause unreliable results if further training is done. Consider using an end-of-epoch checkpoint or enabling fault-tolerant training: https://pytorch-lightning.readthedocs.io/en/stable/advanced/fault_tolerant_training.html
      "You're resuming from a checkpoint that ended before the epoch ended. This can cause unreliable"
    
Epoch 275:   0% 0/342 [00:00<00:00, -72498.33it/s, loss=nan, v_num=3]  [NeMo W 2022-05-09 13:40:53 nemo_logging:349] /usr/local/lib/python3.7/dist-packages/numba/np/ufunc/parallel.py:363: NumbaWarning: The TBB threading layer requires TBB version 2019.5 or later i.e., TBB_INTERFACE_VERSION >= 11005. Found TBB_INTERFACE_VERSION = 9107. The TBB threading layer is disabled.
      warnings.warn(problem)
    
Epoch 275:  12% 40/342 [02:20<-1:58:31, -3.37it/s, loss=2.62, v_num=3]Error executing job with overrides: ['train_dataset=./train.json', 'validation_datasets=./val.json', 'sup_data_path=./fastpitch_sup_data', 'phoneme_dict_path=tts_dataset_files/cmudict-0.7b_nv22.01', 'heteronyms_path=tts_dataset_files/heteronyms-030921', 'whitelist_path=tts_dataset_files/lj_speech.tsv', 'exp_manager.exp_dir=/content/drive/MyDrive/TTS/CHECKPOINTS/checkpoints_dataset_mix_92-1_7hours_2_plus15_16', '+init_from_nemo_model=./tts_en_fastpitch_align.nemo', '+trainer.max_steps=430961', '~trainer.max_epochs', 'trainer.check_val_every_n_epoch=25', 'model.train_ds.dataloader_params.batch_size=24', 'model.validation_ds.dataloader_params.batch_size=24', 'model.n_speakers=2', 'model.pitch_mean=121.9', 'model.pitch_std=20', 'model.pitch_fmin=65', 'model.pitch_fmax=512', 'model.optim.lr=2e-4', '~model.optim.sched', 'model.optim.name=adam', 'trainer.devices=1', 'trainer.strategy=null', '+model.text_tokenizer.add_blank_at=true']
Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 721, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 809, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 1234, in _run
    results = self._run_stage()
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 1321, in _run_stage
    return self._run_train()
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 1351, in _run_train
    self.fit_loop.run()
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/fit_loop.py", line 268, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 208, in advance
    batch_output = self.batch_loop.run(batch, batch_idx)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 88, in advance
    outputs = self.optimizer_loop.run(split_batch, optimizers, batch_idx)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 207, in advance
    self.optimizer_idx,
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 256, in _run_optimization
    self._optimizer_step(optimizer, opt_idx, batch_idx, closure)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 378, in _optimizer_step
    using_lbfgs=is_lbfgs,
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 1593, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/core/lightning.py", line 1644, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/core/optimizer.py", line 168, in step
    step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/strategies/strategy.py", line 193, in optimizer_step
    return self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/plugins/precision/native_amp.py", line 85, in optimizer_step
    closure_result = closure()
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 148, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 143, in closure
    self._backward_fn(step_output.closure_loss)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 311, in backward_fn
    self.trainer._call_strategy_hook("backward", loss, optimizer, opt_idx)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 1763, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/strategies/strategy.py", line 168, in backward
    self.precision_plugin.backward(self.lightning_module, closure_loss, *args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 80, in backward
    model.backward(closure_loss, optimizer, *args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/core/lightning.py", line 1389, in backward
    loss.backward(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/_tensor.py", line 363, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/usr/local/lib/python3.7/dist-packages/torch/autograd/__init__.py", line 175, in backward
    allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass
RuntimeError: cuDNN error: CUDNN_STATUS_INTERNAL_ERROR
You can try to repro this exception using the following code snippet. If that doesn't trigger the error, please include your original repro script when reporting this issue.

import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.allow_tf32 = True
data = torch.randn([24, 1, 1, 119], dtype=torch.half, device='cuda', requires_grad=True)
net = torch.nn.Conv2d(1, 384, kernel_size=[1, 3], padding=[0, 1], stride=[1, 1], dilation=[1, 1], groups=1)
net = net.cuda().half()
out = net(data)
out.backward(torch.randn_like(out))
torch.cuda.synchronize()

ConvolutionParams 
    data_type = CUDNN_DATA_HALF
    padding = [0, 1, 0]
    stride = [1, 1, 0]
    dilation = [1, 1, 0]
    groups = 1
    deterministic = false
    allow_tf32 = true
input: TensorDescriptor 0xebe7090
    type = CUDNN_DATA_HALF
    nbDims = 4
    dimA = 24, 1, 1, 119, 
    strideA = 119, 119, 119, 1, 
output: TensorDescriptor 0xfeaabcd0
    type = CUDNN_DATA_HALF
    nbDims = 4
    dimA = 24, 384, 1, 119, 
    strideA = 45696, 119, 119, 1, 
weight: FilterDescriptor 0x1b618ba0
    type = CUDNN_DATA_HALF
    tensor_format = CUDNN_TENSOR_NCHW
    nbDims = 4
    dimA = 384, 1, 1, 3, 
Pointer addresses: 
    input: 0x7f6abdaa4200
    output: 0x7f6b2b04a000
    weight: 0x7f6a1adff600


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "fastpitch_finetune.py", line 37, in main
    trainer.fit(model)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 769, in fit
    self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 736, in _call_and_handle_interrupt
    self._teardown()
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 1298, in _teardown
    self.strategy.teardown()
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/strategies/single_device.py", line 93, in teardown
    super().teardown()
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/strategies/strategy.py", line 444, in teardown
    optimizers_to_device(self.optimizers, torch.device("cpu"))
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/optimizer.py", line 27, in optimizers_to_device
    optimizer_to_device(opt, device)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/optimizer.py", line 33, in optimizer_to_device
    optimizer.state[p] = apply_to_collection(v, torch.Tensor, move_data_to_device, device)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/apply_func.py", line 108, in apply_to_collection
    v, dtype, function, *args, wrong_dtype=wrong_dtype, include_none=include_none, **kwargs
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/apply_func.py", line 99, in apply_to_collection
    return function(data, *args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/apply_func.py", line 354, in move_data_to_device
    return apply_to_collection(batch, dtype=dtype, function=batch_to)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/apply_func.py", line 99, in apply_to_collection
    return function(data, *args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/apply_func.py", line 347, in batch_to
    data_output = data.to(device, **kwargs)
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
```

And this is the command:

```
!(python fastpitch_finetune.py --config-name=fastpitch_align_v1.05.yaml \
  train_dataset=./train.json \
  validation_datasets=./val.json \
  sup_data_path=./fastpitch_sup_data \
  phoneme_dict_path=tts_dataset_files/cmudict-0.7b_nv22.01 \
  heteronyms_path=tts_dataset_files/heteronyms-030921 \
  whitelist_path=tts_dataset_files/lj_speech.tsv \
  exp_manager.exp_dir=$OUTPUT_CHEKPOINTS \
  +init_from_nemo_model=./tts_en_fastpitch_align.nemo \
  +trainer.max_steps=$max_steps ~trainer.max_epochs \
  trainer.check_val_every_n_epoch=25 \
  model.train_ds.dataloader_params.batch_size=24 model.validation_ds.dataloader_params.batch_size=24 \
  model.n_speakers=2 model.pitch_mean=121.9 model.pitch_std=20 \
  model.pitch_fmin=65 model.pitch_fmax=512 model.optim.lr=2e-4 \
  ~model.optim.sched model.optim.name=adam trainer.devices=1 trainer.strategy=null \
  +model.text_tokenizer.add_blank_at=true \
)
```

It is worth noting that this didn't happen with version 1.6.x. Also, it seems that now training consumes more memory than before, as I had to decrease batch_size in order to keep training.

Thanks,