095 test

.jenkins/download_data.py

@@ -105,6 +105,13 @@ def download_lenet_mnist() -> None:
                          sha256="cb5f8e578aef96d5c1a2cc5695e1aa9bbf4d0fe00d25760eeebaaac6ebc2edcb",
                          )
 
+def download_gpu_quantization_torchao() -> None:
+    # Download SAM model checkpoint for prototype_source/gpu_quantization_torchao_tutorial.py
+    download_url_to_file("https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
+                         prefix=PROTOTYPE_DATA_DIR,
+                         dst="sam_vit_h_4b8939.pth",
+                         sha256="a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e",
+                         )
 
 def main() -> None:
     DATA_DIR.mkdir(exist_ok=True)
@@ -122,7 +129,8 @@ def main() -> None:
         download_dcgan_data()
     if FILES_TO_RUN is None or "fgsm_tutorial" in FILES_TO_RUN:
         download_lenet_mnist()
-
+    if FILES_TO_RUN is None or "gpu_quantization_torchao_tutorial" in FILES_TO_RUN:
+        download_gpu_quantization_torchao()
 
 if __name__ == "__main__":
     main()

.jenkins/metadata.json

@@ -33,5 +33,8 @@
   },
   "intermediate_source/scaled_dot_product_attention_tutorial.py": {
     "needs": "linux.g5.4xlarge.nvidia.gpu"
+  },
+  "prototype_source/gpu_quantization_torchao_tutorial.py": {
+    "needs": "linux.16xlarge.nvidia.gpu"
   }
 }

prototype_source/gpu_quantization_torchao_tutorial.py

@@ -0,0 +1,6 @@
+"""test"""
+import torch
+import time
+
+for i in range(100):
+    a = torch.randn(100,100)

requirements.txt

@@ -6,8 +6,8 @@ sphinx-gallery==0.11.1
 sphinx_design
 docutils==0.16
 sphinx-copybutton
-tqdm
-numpy
+tqdm==4.66.1
+numpy==1.24.4
 matplotlib
 librosa
 torch
@@ -53,6 +53,7 @@ scipy==1.11.1
 numba==0.57.1
 pillow==10.2.0
 wget
+gym==0.26.2
 gym-super-mario-bros==7.4.0
 pyopengl
 gymnasium[mujoco]==0.27.0
@@ -61,3 +62,5 @@ iopath
 pygame==2.1.2
 pycocotools
 semilearn==0.3.2
+torchao==0.0.3
+segment_anything==1.0