Created
June 24, 2024 03:23
-
-
Save tiandiao123/725ad1305268dcf9b8d2cce78ee8203e to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// 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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment