Running Two Moons Notebook with PyTorch Cuda setup yields cuda error

[sjedhoff]

The error seems to originate from the Consistency Model.

Here is the full stack trace:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File <timed exec>:1

File [~\code\python\bayesflow\examples\..\bayesflow\approximators\continuous_approximator.py:114](http://localhost:8888/lab/tree/examples/~/code/python/bayesflow/bayesflow/approximators/continuous_approximator.py#line=113), in ContinuousApproximator.fit(self, *args, **kwargs)
    113 def fit(self, *args, **kwargs):
--> 114     return super().fit(*args, **kwargs, adapter=self.adapter)

File [~\code\python\bayesflow\examples\..\bayesflow\approximators\approximator.py:82](http://localhost:8888/lab/tree/examples/~/code/python/bayesflow/bayesflow/approximators/approximator.py#line=81), in Approximator.fit(self, dataset, simulator, **kwargs)
     80     mock_data = dataset[0]
     81     mock_data = keras.tree.map_structure(keras.ops.convert_to_tensor, mock_data)
---> 82     self.build_from_data(mock_data)
     84 return super().fit(dataset=dataset, **kwargs)

File [~\code\python\bayesflow\examples\..\bayesflow\approximators\approximator.py:23](http://localhost:8888/lab/tree/examples/~/code/python/bayesflow/bayesflow/approximators/approximator.py#line=22), in Approximator.build_from_data(self, data)
     22 def build_from_data(self, data: dict[str, any]) -> None:
---> 23     self.compute_metrics(**data, stage="training")
     24     self.built = True

File [~\code\python\bayesflow\examples\..\bayesflow\approximators\continuous_approximator.py:100](http://localhost:8888/lab/tree/examples/~/code/python/bayesflow/bayesflow/approximators/continuous_approximator.py#line=99), in ContinuousApproximator.compute_metrics(self, inference_variables, inference_conditions, summary_variables, stage)
     97     else:
     98         inference_conditions = keras.ops.concatenate([inference_conditions, summary_outputs], axis=-1)
--> 100 inference_metrics = self.inference_network.compute_metrics(
    101     inference_variables, conditions=inference_conditions, stage=stage
    102 )
    104 loss = inference_metrics.get("loss", keras.ops.zeros(())) + summary_metrics.get("loss", keras.ops.zeros(()))
    106 inference_metrics = {f"{key}/inference_{key}": value for key, value in inference_metrics.items()}

File [~\code\python\bayesflow\examples\..\bayesflow\networks\consistency_models\consistency_model.py:252](http://localhost:8888/lab/tree/examples/~/code/python/bayesflow/bayesflow/networks/consistency_models/consistency_model.py#line=251), in ConsistencyModel.compute_metrics(self, x, conditions, stage)
    251 def compute_metrics(self, x: Tensor, conditions: Tensor = None, stage: str = "training") -> dict[str, Tensor]:
--> 252     base_metrics = super().compute_metrics(x, conditions=conditions, stage=stage)
    254     # The discretization schedule requires the number of passed training steps.
    255     # To be independent of external information, we track it here.
    256     if stage == "training":

File [~\code\python\bayesflow\examples\..\bayesflow\networks\inference_network.py:45](http://localhost:8888/lab/tree/examples/~/code/python/bayesflow/bayesflow/networks/inference_network.py#line=44), in InferenceNetwork.compute_metrics(self, x, conditions, stage)
     43     xz_shape = keras.ops.shape(x)
     44     conditions_shape = None if conditions is None else keras.ops.shape(conditions)
---> 45     self.build(xz_shape, conditions_shape=conditions_shape)
     47 metrics = {}
     49 if stage != "training" and any(self.metrics):
     50     # compute sample-based metrics

File [~\AppData\Local\miniconda3\envs\bf\Lib\site-packages\keras\src\layers\layer.py:225](http://localhost:8888/lab/tree/examples/~/AppData/Local/miniconda3/envs/bf/Lib/site-packages/keras/src/layers/layer.py#line=224), in Layer.__new__.<locals>.build_wrapper(*args, **kwargs)
    223 with obj._open_name_scope():
    224     obj._path = current_path()
--> 225     original_build_method(*args, **kwargs)
    226 # Record build config.
    227 signature = inspect.signature(original_build_method)

File [~\code\python\bayesflow\examples\..\bayesflow\networks\consistency_models\consistency_model.py:160](http://localhost:8888/lab/tree/examples/~/code/python/bayesflow/bayesflow/networks/consistency_models/consistency_model.py#line=159), in ConsistencyModel.build(self, xz_shape, conditions_shape)
    158 for i, n in enumerate(unique_n):
    159     disc = self._discretize_time(n)
--> 160     discretized_times[i, : len(disc)] = disc
    161     discretization_map[n] = i
    162 # Finally, we convert the vectors to tensors

File [~\AppData\Local\miniconda3\envs\bf\Lib\site-packages\torch\_tensor.py:1151](http://localhost:8888/lab/tree/examples/~/AppData/Local/miniconda3/envs/bf/Lib/site-packages/torch/_tensor.py#line=1150), in Tensor.__array__(self, dtype)
   1149     return self.numpy()
   1150 else:
-> 1151     return self.numpy().astype(dtype, copy=False)

TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

[vpratz]

Thanks for reporting! I think the fix in #232 will solve this specific error, I will later test whether there are other CUDA-specific errors in that code.

[vpratz]

I have tested the notebook on a GPU, this problem seems to be resolved. There is still a problem regarding the data adapter, but I will open a different issue for this.