diff --git a/build/context.go b/build/context.go
index d1a1ec5..b537047 100644
--- a/build/context.go
+++ b/build/context.go
@@ -35,7 +35,7 @@ func NewContext() *Context {
 // Package sets the package the generated file will belong to. Required to be able to reference types in the package.
 func (c *Context) Package(path string) {
 	cfg := &packages.Config{
-		Mode: packages.NeedTypes | packages.NeedDeps,
+		Mode: packages.NeedTypes | packages.NeedDeps | packages.NeedImports,
 	}
 	pkgs, err := packages.Load(cfg, path)
 	if err != nil {