diff --git a/__tests__/context.test.ts b/__tests__/context.test.ts
index 63ed0ae..cb6f398 100644
--- a/__tests__/context.test.ts
+++ b/__tests__/context.test.ts
@@ -10,8 +10,8 @@ import {Docker} from '@docker/actions-toolkit/lib/docker/docker';
 import {GitHub} from '@docker/actions-toolkit/lib/github';
 import {Toolkit} from '@docker/actions-toolkit/lib/toolkit';
 
-import {BakeDefinition} from '@docker/actions-toolkit/lib/types/bake';
-import {BuilderInfo} from '@docker/actions-toolkit/lib/types/builder';
+import {BakeDefinition} from '@docker/actions-toolkit/lib/types/buildx/bake';
+import {BuilderInfo} from '@docker/actions-toolkit/lib/types/buildx/builder';
 import {GitHubRepo} from '@docker/actions-toolkit/lib/types/github';
 
 import * as context from '../src/context';
diff --git a/src/context.ts b/src/context.ts
index 1db403e..160d0ce 100644
--- a/src/context.ts
+++ b/src/context.ts
@@ -8,7 +8,7 @@ import {GitHub} from '@docker/actions-toolkit/lib/github';
 import {Toolkit} from '@docker/actions-toolkit/lib/toolkit';
 import {Util} from '@docker/actions-toolkit/lib/util';
 
-import {BakeDefinition} from '@docker/actions-toolkit/lib/types/bake';
+import {BakeDefinition} from '@docker/actions-toolkit/lib/types/buildx/bake';
 
 export interface Inputs {
   builder: string;
diff --git a/src/main.ts b/src/main.ts
index 53d941d..48aa7c9 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -9,8 +9,8 @@ import {Exec} from '@docker/actions-toolkit/lib/exec';
 import {GitHub} from '@docker/actions-toolkit/lib/github';
 import {Toolkit} from '@docker/actions-toolkit/lib/toolkit';
 
-import {BakeDefinition} from '@docker/actions-toolkit/lib/types/bake';
-import {ConfigFile} from '@docker/actions-toolkit/lib/types/docker';
+import {BakeDefinition} from '@docker/actions-toolkit/lib/types/buildx/bake';
+import {ConfigFile} from '@docker/actions-toolkit/lib/types/docker/docker';
 
 import * as context from './context';
 import * as stateHelper from './state-helper';