diff --git a/simple_einet/layers/distributions/abstract_leaf.py b/simple_einet/layers/distributions/abstract_leaf.py
index 3844a3f..6b1edc4 100644
--- a/simple_einet/layers/distributions/abstract_leaf.py
+++ b/simple_einet/layers/distributions/abstract_leaf.py
@@ -1,6 +1,7 @@
 import logging
 from abc import ABC, abstractmethod
 from typing import List
+import numpy as np
 
 import torch
 from torch import distributions as dist, nn