[Bug]: Inconsistent Use of CUDA Devices When Using GPU with notears

[blacksnail789521]

Contact Details

[email protected]

Short description of the problem here.

When using GPU support via from causalnex.structure.pytorch.notears import from_pandas, an error occurs due to some nn.Module components not residing on CUDA devices. The problematic code is:

x = nn.LayerNorm(
    output_dim,
    eps=self.nonlinear_clamp,
    elementwise_affine=True,
)(x)

Link to the code.

To reproduce the error, set hidden_layer_units to [5].

A potential fix might be to append .to(self.device) after LayerNorm, like this:

x = nn.LayerNorm(
    output_dim,
    eps=self.nonlinear_clamp,
    elementwise_affine=True,
).to(self.device)(x)

CausalNex Version

0.12.1

Python Version

3.10.12

Relevant code snippet

x = nn.LayerNorm(
                output_dim,
                eps=self.nonlinear_clamp,
                elementwise_affine=True,
            )(x)

Relevant log output

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument weight in method wrapper_CUDA__native_layer_norm)

Code of Conduct

• I agree to follow this project's Code of Conduct