Fix assertions and tensor device handling in Sinusoidal_embedding class

model.py

@@ -14,13 +14,13 @@ alphas_cumprod = torch.cumprod(alphas, dim=0)
 class Sinusoidal_embedding(torch.nn.Module):
     def __init__(self, dim:int):
         super().__init__()
-        assert self.dim % 2 == 0 , 'Embeddings must be divisble by 2'
-        dim = self.dim 
+        assert dim % 2 == 0 , 'Embeddings must be divisble by 2'
+        self.dim = dim 
 
     def forward(self, time_stamps:torch.Tensor) -> torch.Tensor :
         half_dim = self.dim // 2
         scale = math.log(10000)/(half_dim -1 )
-        freq = torch.exp(torch.arange(half_dim, dtype=torch.float32)* -scale)
+        freq = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=time_stamps.device)* -scale)
         embeddings = time_stamps[:, None] * freq[None, :]
         embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
         return embeddings
@@ -29,7 +29,7 @@ class Sinusoidal_embedding(torch.nn.Module):
 class Time_embeddings(torch.nn.Module):
     def __init__(self, in_dim:int, out_dim:int):
         super().__init__
-        self.Sinusoidal_waves = Sinusoidal_embedding(in_dim)
+        self.Sinusoidal_waves = Sinusoidal_embedding(dim=in_dim)
 
         self.mlp = nn.Sequential(
             nn.Linear(in_dim, out_dim),