Very slow training with GPVAREstimator

[marcotet]

I have been having issues with training GPVAREstimator using GPU, so I had been testing it with a very small test case, and even then it takes multiple minutes to evaluate the model and to do the back-propagation, on either CPU or GPU.

I have been testing this using the solar-nips dataset, restricted to only 10 time-series, and the following parameters for GPVAREstimator:

estimator = GPVAREstimator(prediction_length=dataset.metadata.prediction_length,
                               freq=dataset.metadata.freq,
                               target_dim=n_timeseries,
                               num_layers=2,
                               num_cells=2,
                               cell_type="gru",
                               rank=1,
                               dropout_rate=0.2,
                               trainer=Trainer(
                                               epochs=1,
                                               num_batches_per_epoch=1,
                                               ctx=context.gpu() if device == "gpu" else context.cpu(),
                               )
)

So my question is whether it is normal that GPVAREstimator is slow even with that few layers and cells, or if there is a possibility that I have done something wrong elsewhere in my environment.