From ba5043ddf8c2ed3c5ea00e8e430c8a3d08b737d1 Mon Sep 17 00:00:00 2001
From: Tarry Singh <tarry.singh@deepkapha.com>
Date: Mon, 1 Dec 2025 14:34:13 +0000
Subject: [PATCH] Switch ZML to use .host_pinned memory on ROCm, addressing
 missing host_unpinned memory error.

---
 zml/platform.zig | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/zml/platform.zig b/zml/platform.zig
index 410d60b..c5872a5 100644
--- a/zml/platform.zig
+++ b/zml/platform.zig
@@ -89,7 +89,9 @@ pub const Platform = struct {
         const memory_target: pjrt.Memory.Kind = switch (memory) {
             .host_unpinned => switch (platform.target) {
                 // Cuda doesn't have host_unpinned.
-                .cuda => .host_pinned,
+                // ROCm doesn't seem to have it either.
+                // TODO(gwenzek): investigate why it was not forced before.
+                .cuda, .rocm => .host_pinned,
                 else => .host_unpinned,
             },
             inline else => |t| t,