diff --git a/s2sa/extract_states.lua b/s2sa/extract_states.lua
index 41374cc..cb67b37 100644
--- a/s2sa/extract_states.lua
+++ b/s2sa/extract_states.lua
@@ -3,9 +3,9 @@ require 'string'
 require 'hdf5'
 require 'nngraph'
 
-require 'models.lua'
-require 'data.lua'
-require 'util.lua'
+require 's2sa.models'
+require 's2sa.data'
+require 's2sa.model_utils'
 
 stringx = require('pl.stringx')