BayesianNetwork.predict throws error when using pytorch backend

[ankurankan]

In [8]: model = get_example_model('cancer')                                             [285/285]
In [8]: model                                   
Out[8]: <pgmpy.models.BayesianNetwork.BayesianNetwork at 0x7efffd538ac0>                                                                         
In [9]: model.nodes()                                                                            
Out[9]: NodeView(('Pollution', 'Smoker', 'Cancer', 'Xray', 'Dyspnoea'))                          
                                                                                                 
In [10]: df = model.simulate(100)                                                                
                                                                                                
WARNING:pgmpy:Probability values don't exactly sum to 1. Differ by: 2.2351741790771484e-08. Adjust
ing values.                                                                                      Generating for node: Dyspnoea: 100%|██████████████████████████████| 5/5 [00:00<00:00, 695.46it/s]
/home/ankur/pgmpy/pgmpy/sampling/base.py:598: FutureWarning: Passing a DataFrame to DataFrame.fro
m_records is deprecated. Use set_index and/or drop to modify the DataFrame instead.              
  df = pd.DataFrame.from_records(samples)                                                        
                                                                                                 
In [12]: df_small = df.drop('Xray', axis=1)                                                      
                                                                                                 
In [13]: model.predict(df_small)   

Error

ValueError: einstein sum subscripts string contains too many subscripts for operand 0