diff --git a/nerfies/configs.py b/nerfies/configs.py
index f349804..079689e 100644
--- a/nerfies/configs.py
+++ b/nerfies/configs.py
@@ -16,7 +16,7 @@
 from typing import Any, Mapping, Optional, Tuple
 
 import dataclasses
-from flax import nn
+from flax.linen import nn
 import gin
 import immutabledict