diff --git a/Dockerfile b/Dockerfile
index 3504e99..8bacee9 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,5 +1,8 @@
-# Use official Python image as base
-FROM python:3.10-slim
+# NVIDIA
+#FROM python:3.10-slim
+
+# AMD
+FROM rocm/pytorch:latest
 
 # Set working directory
 WORKDIR /app