solved test problem

Jorgedavyd · Jun 10, 2024 · 78e22e4 · 78e22e4
1 parent f16ea7b
commit 78e22e4
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/lightorch/nn/functional.py b/lightorch/nn/functional.py
@@ -2,10 +2,9 @@
 from torch import nn, Tensor
 from typing import Optional, Union, Tuple, Callable, List
 import torch.nn.functional as F
-from lightning.pytorch import LightningModule
 from einops import rearrange
 from torch.fft import fftn
-from .utils import FeatureExtractor
+from .utils import FeatureExtractor2D
 
 
 def _fourierconvNd(n: int, x: Tensor, weight: Tensor, bias: Union[Tensor,None]) -> Tensor:
@@ -358,7 +357,7 @@ def style_loss(
     input: Tensor,
     target: Tensor,
     F_p: Tensor,
-    feature_extractor: FeatureExtractor = None,
+    feature_extractor: FeatureExtractor2D = None,
 ) -> Tensor:
     if feature_extractor is not None:
         phi_input: Tensor = feature_extractor(input)
@@ -377,7 +376,7 @@ def perceptual_loss(
     input: Tensor,
     target: Tensor,
     N_phi_p: Tensor,
-    feature_extractor: FeatureExtractor = None,
+    feature_extractor: FeatureExtractor2D = None,
 ) -> Tensor:
     if feature_extractor is not None:
         phi_input: Tensor = feature_extractor(input)