Skip to content

Latest commit

 

History

History
19 lines (16 loc) · 681 Bytes

File metadata and controls

19 lines (16 loc) · 681 Bytes
title How do I make W&B Launch work with Tensorflow on GPU?
description Configure a custom accelerator base image so W&B Launch jobs run TensorFlow with GPU support.

For TensorFlow jobs using GPUs, specify a custom base image for the container build. This ensures proper GPU utilization during runs. Add an image tag under the builder.accelerator.base_image key in the resource configuration. For example:

{
    "gpus": "all",
    "builder": {
        "accelerator": {
            "base_image": "tensorflow/tensorflow:latest-gpu"
        }
    }
}

In versions prior to W&B 0.15.6, use cuda instead of accelerator as the parent key for base_image.