🔨 Modified activation in TimeEmbed function

Author: marianovitasari20

ice_station_zebra/models/common/timeembed.py

@@ -1,14 +1,20 @@
 import torch.nn as nn
 from torch import Tensor
 
+from .activations import get_activation
+
 
 class TimeEmbed(nn.Module):
-    def __init__(self, dim: int = 256) -> None:
+    def __init__(self, dim: int = 256, 
+                 activation: str = "SiLU",) -> None:
         super().__init__()
 
+        def act():
+            return get_activation(activation)
+
         self.model = nn.Sequential(
             nn.Linear(dim, dim * 4),
-            nn.SiLU(),
+            act(),
             nn.Linear(dim * 4, dim),
         )