UnboundLocalError when loading pre-trained VAE model: Variable 'T' not initialized

[sd68515]

Hello,
When attempting to load a pre-trained VAE model (encoder_vae.pt/decoder_vae.pt) in a new environment, initialization fails with:

adata=sc.read("data_dir")

adata
AnnData object with n_obs × n_vars = 54941 × 2000
    obs: 'Sample.ID', 'Age', 'Tissue.ID', 'Sample.Type', 'Assay', 'Batch', 'seurat_clusters', 'RNA.Counts', 'RNA.Features', 'Percent.MT', 'Percent.Ribo', 'Cell.Barcode', 'DF_pANN', 'DF_classification', 'DF_pANN_quantile', 'Spliced.Counts', 'Spliced.Features', 'Unspliced.Counts', 'Unspliced.Features', 'Ambiguous.Counts', 'Ambiguous.Features', 'cell_type', 'n_counts', 'n_genes', 'initial_size_spliced', 'initial_size_unspliced', 'initial_size', 'clusters', 'vae_time', 'vae_std_t', 'vae_t0', 'vae_velocity_self_transition', 'root_cells', 'end_points', 'vae_velocity_pseudotime', 'vae_velocity_length', 'vae_velocity_confidence', 'vae_velocity_confidence_transition', 'velocity_cosine_sim'
    var: 'gene_count_corr', 'means', 'dispersions', 'dispersions_norm', 'highly_variable', 'init_mode', 'w_init', 'vae_alpha', 'vae_beta', 'vae_gamma', 'vae_ton', 'vae_scaling', 'vae_sigma_u', 'vae_sigma_s'
    uns: 'cell_type_colors', 'leiden', 'neighbors', 'pca', 'vae_run_time', 'vae_test_idx', 'vae_train_idx', 'vae_velocity_graph', 'vae_velocity_graph_neg', 'vae_velocity_graph_uncertainties', 'vae_velocity_params'
    obsm: 'X_pca', 'X_umap', 'vae_std_z', 'vae_velocity_umap', 'vae_z'
    varm: 'PCs', 'vae_mode'
    layers: 'Ms', 'Mu', 'spliced', 'unspliced', 'vae_rho', 'vae_s0', 'vae_shat', 'vae_u0', 'vae_uhat', 'vae_velocity', 'vae_velocity_u'
    obsp: 'connectivities', 'distances'

vae = vv.VAE(adata, tmax=20, dim_z=5, checkpoints=[os.path.join(model_path, f"encoder_vae.pt"), os.path.join(model_path,f"decoder_vae.pt")])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/data_d/Velocity/LXiao/VeloVAE/velovae/model/vae.py", line 625, in __init__
    self.decoder = decoder(adata,
                   ^^^^^^^^^^^^^^
  File "/data_d/Velocity/LXiao/VeloVAE/velovae/model/vae.py", line 293, in __init__
    dyn_mask = (T > tmax*0.01) & (np.abs(T-toff) > tmax*0.01)
                ^
UnboundLocalError: cannot access local variable 'T' where it is not associated with a value

The adata object was saved after running the standard workflow:

vv.preprocess(adata, 2000, compute_umap=False, perform_clustering=True, npc=30)

vae = vv.VAE(adata, tmax=20, dim_z=5, device='cuda:1', reverse_gene_mode=False, full_vb=False, discrete=False)

vae.train(adata, figure_path=figure_path, embed='umap')

res, res_type = vv.post_analysis(adata, id_pre, methods=['VeloVAE'], keys=['vae'], cluster_key=clusterkey, dpi=600, compute_metrics=False, raw_count=False, figure_path=figure_path, n_jobs=32, save=data_dir)

How is T supposed to be initialized when loading a pre-trained model?
And why does this error occur only for some datasets despite identical pipelines?

Your support on this matter would be greatly appreciated.
Best regards,