fix torch import for flavors (#1129)

unit8co · Aug 7, 2022 · eb18103 · eb18103
1 parent 4e5f1e6
commit eb18103
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/darts/tests/models/forecasting/test_transformer_model.py b/darts/tests/models/forecasting/test_transformer_model.py
@@ -2,7 +2,6 @@
 import tempfile
 
 import pandas as pd
-import torch.nn as nn
 
 from darts import TimeSeries
 from darts.logging import get_logger
@@ -12,6 +11,8 @@
 logger = get_logger(__name__)
 
 try:
+    import torch.nn as nn
+
     from darts.models.components.transformer import (
         CustomFeedForwardDecoderLayer,
         CustomFeedForwardEncoderLayer,