Customer churn model training error

[gdprtcl]

I'm working on MLOps with the example of customer_churn.
During HyperParameter Tuning stage, I'm getting the following error:
User program failed with RuntimeError: expand(torch.FloatTensor{[2, 64]}, size=[64]): the number of sizes provided (1) must be greater or equal to the number of dimensions in the tensor (2)

The experiment failed. Finalizing run...
Cleaning up all outstanding Run operations, waiting 300.0 seconds
2 items cleaning up...
Cleanup took 0.13342905044555664 seconds
Starting the daemon thread to refresh tokens in background for process with pid = 115
Traceback (most recent call last):
File "svdkl_entry.py", line 61, in
trainer.fit(train_dataloader)
File "/mnt/batch/tasks/shared/LS_root/jobs/mlops-aml-ws/azureml/hd_88232b3c-5bdb-4f2f-8dad-6edb3674a6ce_3/mounts/workspaceblobstore/azureml/HD_88232b3c-5bdb-4f2f-8dad-6edb3674a6ce_3/trainer.py", line 87, in fit
output = self.model(data)
File "/azureml-envs/azureml_8a75b4939559f760114d357d5e253d59/lib/python3.6/site-packages/gpytorch/module.py", line 24, in call
outputs = self.forward(*inputs, **kwargs)
File "/mnt/batch/tasks/shared/LS_root/jobs/mlops-aml-ws/azureml/hd_88232b3c-5bdb-4f2f-8dad-6edb3674a6ce_3/mounts/workspaceblobstore/azureml/HD_88232b3c-5bdb-4f2f-8dad-6edb3674a6ce_3/svdkl.py", line 145, in forward
res = self.gp_layer(features)
File "/azureml-envs/azureml_8a75b4939559f760114d357d5e253d59/lib/python3.6/site-packages/gpytorch/models/approximate_gp.py", line 81, in call
return self.variational_strategy(inputs, prior=prior)
File "/azureml-envs/azureml_8a75b4939559f760114d357d5e253d59/lib/python3.6/site-packages/gpytorch/variational/_variational_strategy.py", line 108, in call
self.variational_distribution.initialize_variational_distribution(prior_dist)
File "/azureml-envs/azureml_8a75b4939559f760114d357d5e253d59/lib/python3.6/site-packages/gpytorch/variational/cholesky_variational_distribution.py", line 50, in initialize_variational_distribution
self.variational_mean.data.copy(prior_dist.mean)
RuntimeError: expand(torch.FloatTensor{[2, 64]}, size=[64]): the number of sizes provided (1) must be greater or equal to the number of dimensions in the tensor (2)