diff --git a/__tests__/context.test.ts b/__tests__/context.test.ts
index 7149b27..720fd21 100644
--- a/__tests__/context.test.ts
+++ b/__tests__/context.test.ts
@@ -6,7 +6,7 @@ import {Buildx} from '@docker/actions-toolkit/lib/buildx/buildx';
 import {Context} from '@docker/actions-toolkit/lib/context';
 import {Docker} from '@docker/actions-toolkit/lib/docker/docker';
 import {Toolkit} from '@docker/actions-toolkit/lib/toolkit';
-import {Node} from '@docker/actions-toolkit/lib/types/builder';
+import {Node} from '@docker/actions-toolkit/lib/types/buildx/builder';
 
 import * as context from '../src/context';
 
diff --git a/src/context.ts b/src/context.ts
index cdaa4c3..587b3d3 100644
--- a/src/context.ts
+++ b/src/context.ts
@@ -1,9 +1,11 @@
 import * as uuid from 'uuid';
 import * as core from '@actions/core';
+
 import {Docker} from '@docker/actions-toolkit/lib/docker/docker';
 import {Util} from '@docker/actions-toolkit/lib/util';
 import {Toolkit} from '@docker/actions-toolkit/lib/toolkit';
-import {Node} from '@docker/actions-toolkit/lib/types/builder';
+
+import {Node} from '@docker/actions-toolkit/lib/types/buildx/builder';
 
 export const builderNodeEnvPrefix = 'BUILDER_NODE';
 
diff --git a/src/main.ts b/src/main.ts
index bd8c597..bdaf861 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -2,13 +2,15 @@ import * as fs from 'fs';
 import * as yaml from 'js-yaml';
 import * as core from '@actions/core';
 import * as actionsToolkit from '@docker/actions-toolkit';
+
 import {Buildx} from '@docker/actions-toolkit/lib/buildx/buildx';
 import {Builder} from '@docker/actions-toolkit/lib/buildx/builder';
 import {Docker} from '@docker/actions-toolkit/lib/docker/docker';
 import {Exec} from '@docker/actions-toolkit/lib/exec';
 import {Toolkit} from '@docker/actions-toolkit/lib/toolkit';
 import {Util} from '@docker/actions-toolkit/lib/util';
-import {Node} from '@docker/actions-toolkit/lib/types/builder';
+
+import {Node} from '@docker/actions-toolkit/lib/types/buildx/builder';
 
 import * as context from './context';
 import * as stateHelper from './state-helper';