cudnn RNN backward can only be called in training mode

[ali-akhavan89]

Hello,

I have put my custom embedding network in training mode, sequence_net.train() , before the training happens, but I still get the following error. This happens when I use 'mdn' during multi-round (right when the non-atomic training starts). I don't have this issue with 'maf' in multi-round training. Everything is on cuda. I would appreciate any thoughts or feedback. Thank you!

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[1], line 109
    106 warnings.filterwarnings("ignore", category=UserWarning)
    107 vensimInputs['device'] = str(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
--> 109 posterior = run_training(vensimInputs, SBI_Inputs, resInputs)  

File D:\Ali\validation_v07\setup_functions\run_training.py:203, in run_training(vensimInputs, SBI_Inputs, resInputs)
    199 sequence_net.train()
    201 print("sequence_net mode before inference.train():", sequence_net.training)  # Should be True
--> 203 density_estimator = inference.train(training_batch_size=SBI_Inputs['training_batch_size'],
    204                                     stop_after_epochs=SBI_Inputs['stop_after_epochs'],
    205                                     learning_rate=SBI_Inputs['learning_rate'],
    206                                     show_train_summary=SBI_Inputs['show_train_summary'])
    207 train_time = round(time.time()-start_time)
    208 print(f"Training time: {train_time//3600}:{(train_time%3600)//60:02}:{train_time%60:02}\n\n")

File ~\.conda\envs\ptgpu\Lib\site-packages\sbi\inference\trainers\npe\npe_c.py:189, in NPE_C.train(self, num_atoms, training_batch_size, learning_rate, validation_fraction, stop_after_epochs, max_num_epochs, clip_max_norm, calibration_kernel, resume_training, force_first_round_loss, discard_prior_samples, use_combined_loss, retrain_from_scratch, show_train_summary, dataloader_kwargs)
    185     if self.use_non_atomic_loss:
    186         # Take care of z-scoring, pre-compute and store prior terms.
    187         self._set_state_for_mog_proposal()
--> 189 return super().train(**kwargs)

File ~\.conda\envs\ptgpu\Lib\site-packages\sbi\inference\trainers\npe\npe_base.py:373, in PosteriorEstimator.train(self, training_batch_size, learning_rate, validation_fraction, stop_after_epochs, max_num_epochs, clip_max_norm, calibration_kernel, resume_training, force_first_round_loss, discard_prior_samples, retrain_from_scratch, show_train_summary, dataloader_kwargs)
    370 train_loss = torch.mean(train_losses)
    371 train_loss_sum += train_losses.sum().item()
--> 373 train_loss.backward()
    374 if clip_max_norm is not None:
    375     clip_grad_norm_(
    376         self._neural_net.parameters(), max_norm=clip_max_norm
    377     )

File ~\.conda\envs\ptgpu\Lib\site-packages\torch\_tensor.py:581, in Tensor.backward(self, gradient, retain_graph, create_graph, inputs)
    571 if has_torch_function_unary(self):
    572     return handle_torch_function(
    573         Tensor.backward,
    574         (self,),
   (...)
    579         inputs=inputs,
    580     )
--> 581 torch.autograd.backward(
    582     self, gradient, retain_graph, create_graph, inputs=inputs
    583 )

File ~\.conda\envs\ptgpu\Lib\site-packages\torch\autograd\__init__.py:347, in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    342     retain_graph = create_graph
    344 # The reason we repeat the same comment below is that
    345 # some Python versions print out the first line of a multi-line function
    346 # calls in the traceback and some print out the last line
--> 347 _engine_run_backward(
    348     tensors,
    349     grad_tensors_,
    350     retain_graph,
    351     create_graph,
    352     inputs,
    353     allow_unreachable=True,
    354     accumulate_grad=True,
    355 )

File ~\.conda\envs\ptgpu\Lib\site-packages\torch\autograd\graph.py:825, in _engine_run_backward(t_outputs, *args, **kwargs)
    823     unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
    824 try:
--> 825     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    826         t_outputs, *args, **kwargs
    827     )  # Calls into the C++ engine to run the backward pass
    828 finally:
    829     if attach_logging_hooks:

RuntimeError: cudnn RNN backward can only be called in training mode