From f9316e359a06ceb19c5f8ed3d5f5c4113b5bda78 Mon Sep 17 00:00:00 2001
From: Dario Coscia <93731561+dario-coscia@users.noreply.github.com>
Date: Thu, 27 Jun 2024 17:36:54 +0200
Subject: [PATCH] gpu fourierembedding (#313)

Adding gpu to fourier embedding
---
 pina/model/layers/embedding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pina/model/layers/embedding.py b/pina/model/layers/embedding.py
index c9f41da..4248136 100644
--- a/pina/model/layers/embedding.py
+++ b/pina/model/layers/embedding.py
@@ -246,7 +246,7 @@ class FourierFeatureEmbedding(torch.nn.Module):
         :rtype: torch.Tensor
         """
         # compute random matrix multiplication
-        out = torch.mm(x, self._matrix)
+        out = torch.mm(x, self._matrix.to(device=x.device, dtype=x.dtype))
         # return embedding
         return torch.cat(
             [torch.cos(2 * torch.pi * out), torch.sin(2 * torch.pi * out)],