Makefile that I currently have:
NVCC=nvcc
RM=rm -rf
PYTHON_HEADER_DIR := $(shell python -c 'from distutils.sysconfig import get_python_inc; print(get_python_inc())')
PYTORCH_INCLUDES := $(shell python -c 'from torch.utils.cpp_extension import include_paths; [print(p) for p in include_paths()]')
PYTORCH_LIBRARIES := $(shell python -c 'from torch.utils.cpp_extension import library_paths; [print(p) for p in library_paths()]')

# CUDA ROOT DIR that contains bin/ lib64/ and include/
# CUDA_DIR := /usr/local/cuda
CUDA_DIR := $(shell python -c 'from torch.utils.cpp_extension import _find_cuda_home; print(_find_cuda_home())')
INCLUDE_DIRS := ./ $(CUDA_DIR)/include
INCLUDE_DIRS += $(PYTHON_HEADER_DIR)
INCLUDE_DIRS += $(PYTORCH_INCLUDES)
COMMON_FLAGS += $(foreach includedir,$(INCLUDE_DIRS),-I$(includedir)) -DTORCH_API_INCLUDE_EXTENSION_H -D_GLIBCXX_USE_CXX11_ABI=0


CUDA_ARCH := -gencode arch=compute_35,code=sm_35 \
#               -gencode arch=compute_35,code=sm_35 \
                -gencode arch=compute_50,code=sm_50 \
                -gencode arch=compute_52,code=sm_52 \

#LIBRARIES += stdc++ cudart c10 caffe2 torch torch_python caffe2_gpu
LIBS= -lopenblas -lpthread -lcudart -lcublas
LIBS += $(PYTORCH_LIBRARIES)

NVCCFLAGS= -std=c++11 -c $(CUDA_ARCH) -ccbin=$(CXX) -Xcompiler -fPIC $(COMMON_FLAGS)
CCFLAGS= -O3 -std=c++11 $(COMMON_FLAGS)


all: build

build: cpu gpu
        $(NVCC) $(NVCCFLAGS) -o test *.o

cpu:
        $(CC) $(CCFLAGS) *.cpp

gpu:
        $(NVCC) $(NVCCFLAGS) *.cu

clean:
        $(RM) test *.o