diff --git a/elasticai/creator/base_modules/layers.py b/elasticai/creator/base_modules/layers.py
index c43cbe7a..de3cf491 100644
--- a/elasticai/creator/base_modules/layers.py
+++ b/elasticai/creator/base_modules/layers.py
@@ -1,10 +1,10 @@
 from typing import Callable, Optional, cast
 
 import torch
+from torch import Tensor
+from torch.nn import Module
 from torch.nn.utils.parametrize import register_parametrization
 
-from elasticai.creator.mlframework import Module, Tensor
-
 from .autograd_functions.binary_quantization import Binarize as BinarizeFn
 
 """Implementation of quantizers and quantized variants of pytorch layers"""
diff --git a/elasticai/creator/nn/conv1d/design.py b/elasticai/creator/nn/conv1d/design.py
index 90700301..de949562 100644
--- a/elasticai/creator/nn/conv1d/design.py
+++ b/elasticai/creator/nn/conv1d/design.py
@@ -7,8 +7,8 @@
 )
 from elasticai.creator.vhdl.design.design import Design
 from elasticai.creator.vhdl.design.ports import Port
-from elasticai.creator.vhdl.designs.rom import Rom
 from elasticai.creator.vhdl.savable import Path
+from elasticai.creator.vhdl.shared_designs.rom import Rom
 
 
 class FPConv1d(Design):
diff --git a/elasticai/creator/nn/linear/design.py b/elasticai/creator/nn/linear/design.py
index 68abab0e..23dbd049 100644
--- a/elasticai/creator/nn/linear/design.py
+++ b/elasticai/creator/nn/linear/design.py
@@ -8,8 +8,8 @@
 )
 from elasticai.creator.vhdl.design.design import Design
 from elasticai.creator.vhdl.design.ports import Port
-from elasticai.creator.vhdl.designs.rom import Rom
 from elasticai.creator.vhdl.savable import Path
+from elasticai.creator.vhdl.shared_designs.rom import Rom
 
 
 class FPLinear(Design):
diff --git a/elasticai/creator/nn/precomputed/fp_precomputed_module.py b/elasticai/creator/nn/precomputed/fp_precomputed_module.py
index 99392ee0..b064b2f7 100644
--- a/elasticai/creator/nn/precomputed/fp_precomputed_module.py
+++ b/elasticai/creator/nn/precomputed/fp_precomputed_module.py
@@ -12,7 +12,7 @@
     FixedPointConfig,
 )
 from elasticai.creator.vhdl.design.design import Design
-from elasticai.creator.vhdl.designs.precomputed_scalar_function import (
+from elasticai.creator.vhdl.shared_designs.precomputed_scalar_function import (
     PrecomputedScalarFunction,
 )
 from elasticai.creator.vhdl.translatable import Translatable
diff --git a/tests/vhdl/designs/rom/test_rom_file.py b/tests/vhdl/designs/rom/test_rom_file.py
index 2a71f9b3..3458fe5d 100644
--- a/tests/vhdl/designs/rom/test_rom_file.py
+++ b/tests/vhdl/designs/rom/test_rom_file.py
@@ -5,7 +5,7 @@
 import pytest
 
 from elasticai.creator.in_memory_path import InMemoryFile, InMemoryPath
-from elasticai.creator.vhdl.designs.rom import Rom
+from elasticai.creator.vhdl.shared_designs.rom import Rom
 
 from .utils import extract_rom_values