[CLEANUP]

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zetascale"
-version = "2.1.6"
+version = "2.1.7"
 description = "Transformers at zeta scales"
 authors = ["Zeta Team <[email protected]>"]
 license = "MIT"

zeta/nn/modules/__init__.py

@@ -188,6 +188,7 @@
 from zeta.nn.modules.ws_conv2d import WSConv2d
 from zeta.nn.modules.yolo import yolo
 from zeta.nn.modules.palo_ldp import PaloLDP
+
 # from zeta.nn.modules.g_shard_moe import (
 # Top1Gate,
 # Top2Gate,

zeta/nn/modules/palo_ldp.py

@@ -1,7 +1,7 @@
-
 from torch import Tensor, nn
 from zeta.utils.log_pytorch_op import log_torch_op
 
+
 class PaloLDP(nn.Module):
  """
  Implementation of the PaloLDP module.
@@ -75,7 +75,6 @@ def forward(self, x: Tensor) -> Tensor:
  x = self.pointwise_conv(x)
  print(x.shape) # torch.Size([2, 1, 4, 4]
 
-
  # Depthwise convolution with 1 stide
  x = self.depthwise_conv(x)
  print(x.shape)
@@ -89,7 +88,7 @@ def forward(self, x: Tensor) -> Tensor:
  print(x.shape)
 
  # Norm
- x = self.norm(x) #+ skip
+ x = self.norm(x)  # + skip
  print(x.shape)
 
  # Depthwise convolution with 2 stide
@@ -109,5 +108,3 @@ def forward(self, x: Tensor) -> Tensor:
  x = nn.LayerNorm(w)(x)
 
  return x
-
-

zeta/utils/__init__.py

@@ -90,5 +90,5 @@
  "check_cuda",
  "VerboseExecution",
  "seek_all_images",
- "log_torch_op"
+ "log_torch_op",
 ]

zeta/utils/log_pytorch_op.py

@@ -13,8 +13,8 @@
  format="<green>{time}</green> <level>{message}</level>",
  backtrace=True,
  diagnose=True,
- enqueue = True,
- catch = True,
+ enqueue=True,
+ catch=True,
 )
 
 
@@ -39,6 +39,7 @@ def log_torch_op(
  Returns:
  function: The decorated function.
  """
+
  def decorator(func):
  @functools.wraps(func)
  def wrapper(*args, **kwargs):