1. System Info

Environment

• PyPOTS Version: 1.0
• PyTorch Version: (installed via dependencies)
• Python Version: 3.11
• Platform: Linux

2. Information

• The official example scripts
• My own created scripts

3. Reproduction

Description

The GPVAE model initialization succeeds, but training fails immediately during the first forward pass. The error occurs in the matern_kernel function when attempting to compute torch.sqrt(length_scale) where length_scale is a Python float instead of a PyTorch Tensor.

Error Traceback

2025-11-08 12:44:29 [ERROR]: ❌ Exception: sqrt(): argument 'input' (position 1) must be Tensor, not float
Error: Training got interrupted. Model was not trained. Please investigate the error printed above.
Traceback (most recent call last):
  File "/home/n66w96/timeseriesimputation/.venv/lib/python3.11/site-packages/pypots/base.py", line 737, in _train_model
    results = self.model(inputs, calc_criterion=True)
  File "/home/n66w96/timeseriesimputation/.venv/lib/python3.11/site-packages/pypots/nn/modules/gpvae/backbone.py", line 160, in forward
    self.prior = self._init_prior(device=X.device)
  File "/home/n66w96/timeseriesimputation/.venv/lib/python3.11/site-packages/pypots/nn/modules/gpvae/backbone.py", line 121, in _init_prior
    kernel_matrices.append(matern_kernel(self.time_length, self.length_scale / 2**i))
  File "/home/n66w96/timeseriesimputation/.venv/lib/python3.11/site-packages/pypots/nn/modules/gpvae/layers.py", line 39, in matern_kernel
    distance_matrix_scaled = distance_matrix / torch.sqrt(length_scale).type(torch.float32)
TypeError: sqrt(): argument 'input' (position 1) must be Tensor, not float

Root Cause

File: /pypots/nn/modules/gpvae/layers.py
Line: 39
Issue: The matern_kernel function receives length_scale as a Python float (from self.length_scale / 2**i) and attempts to pass it directly to torch.sqrt(), which only accepts Tensor objects.

Minimal Reproduction Code

from pypots.imputation import GPVAE
import numpy as np

np.random.seed(42)
data = np.random.randn(10, 96, 49)
data_missing = data.copy()
data_missing[np.random.rand(*data_missing.shape) < 0.2] = np.nan

train_set = {'X': data_missing}
val_set = {'X': data_missing, 'X_ori': data}

model = GPVAE(
    n_steps=96,
    n_features=49,
    latent_size=45,
    encoder_sizes=(128, 128),
    decoder_sizes=(512, 512),
    kernel='matern',
    beta=0.849,
    M=2,
    K=1,
    sigma=1.01,
    length_scale=8.0,
    kernel_scales=2,
    window_size=9,
    batch_size=64,
    epochs=1
)

# Model creation succeeds
print('Model created successfully')

# Training fails with TypeError
model.fit(train_set, val_set)

4. Expected behavior

Expected Behavior

The model should train successfully without type errors. The length_scale parameter should be properly converted to a Tensor before being used in tensor operations.

5. Your contribution

Proposed Fix

In /pypots/nn/modules/gpvae/layers.py line 39, convert length_scale to a tensor:

Current code:

distance_matrix_scaled = distance_matrix / torch.sqrt(length_scale).type(torch.float32)

Proposed fix:

length_scale_tensor = torch.tensor(length_scale, dtype=torch.float32, device=distance_matrix.device)
distance_matrix_scaled = distance_matrix / torch.sqrt(length_scale_tensor)

Alternatively, ensure length_scale is already a Tensor when passed to the function.