Memory_problems.jl

using Flux

const nn =  gpu(Chain(
            Conv((5, 5), 3=>6, relu),
            MaxPool((2, 2)),
            Conv((5, 5), 6=>16, relu),
            MaxPool((2, 2)),
            flatten,
            Dense(256, 120, relu), 
            Dense(120, 84, relu), 
            Dense(84, 4)
          ))

function inference(imgs)
    out = nn(imgs)
    maximum(cpu(out))
end

imgs = gpu(randn(28,28,3,1000))
N = 100

# Example 1, Threaded execution, crashes after
#  about 10 runs with 11GB GPU Memory
res = Any[nothing for i in 1:N]
@sync for i in 1:N
    res[i] = Threads.@spawn inference(imgs)
end
s = sum(fetch.(res))


# Example 2, Serially threaded, crashes after
#  about 10 runs with 11GB GPU Memory
for j in 1:10
    res = Any[nothing for i in 1:N]
    for i in 1:N
        res[i] = fetch(Threads.@spawn inference(imgs))
    end
    s = sum(res)
    println(s)
end