Dockerfile for GPU and non-GPU versions

I have a project that I maintain using Dockerfile. It is a data science project so and the docker can be used in many environments, local, in the cloud, with or without GPU.

The difference between the GPU and the non-GPU version is very small, just this:

FROM ubuntu:18.04  -> non-GPU
FROM nvidia/cuda   -> GPU 

What is the best way to handle this scenario?

P.S: I am using docker-compose to build the docker

Source: Docker Questions