folivo/PINA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gpu fourierembedding (#313)

Browse Source 
Adding gpu to fourier embedding
This commit is contained in:
Dario Coscia
2024-06-27 17:36:54 +02:00
committed by GitHub
parent ba79984be2
commit f9316e359a
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
pina/model/layers/embedding.py
View File

@@ -246,7 +246,7 @@ class FourierFeatureEmbedding(torch.nn.Module):
:rtype: torch.Tensor :rtype: torch.Tensor
""" """
# compute random matrix multiplication # 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 embedding
return torch.cat( return torch.cat(
[torch.cos(2 * torch.pi * out), torch.sin(2 * torch.pi * out)], [torch.cos(2 * torch.pi * out), torch.sin(2 * torch.pi * out)],