concurrent_queue.h

// concurrent_queue.h
#ifndef CONCURRENT_QUEUE_H
#define CONCURRENT_QUEUE_H

#include <torch/torch.h>
#include <queue>
#include <mutex>
#include <condition_variable>

template <typename T>
class ConcurrentQueue {
public:
    void push(T value);
    bool try_pop(T& value);
    void wait_and_pop(T& value);
    bool empty() const;

private:
    std::queue<T> queue;
    mutable std::mutex mutex;
    std::condition_variable cv;
};

void producer(ConcurrentQueue<torch::Tensor>& queue);
void consumer(ConcurrentQueue<torch::Tensor>& queue);

#endif // CONCURRENT_QUEUE_TORCH_H