mirror of
https://github.com/ml-explore/mlx.git
synced 2025-11-04 10:38:10 +08:00
25 lines
517 B
C++
25 lines
517 B
C++
// Copyright © 2023 Apple Inc.
|
|
|
|
#define DOCTEST_CONFIG_IMPLEMENT
|
|
#include "doctest/doctest.h"
|
|
|
|
#include <cstdlib>
|
|
|
|
#include "mlx/mlx.h"
|
|
|
|
using namespace mlx::core;
|
|
|
|
int main(int argc, char** argv) {
|
|
doctest::Context context;
|
|
|
|
const char* device = std::getenv("DEVICE");
|
|
if (device != nullptr && std::string(device) == "cpu") {
|
|
set_default_device(Device::cpu);
|
|
} else if (metal::is_available()) {
|
|
set_default_device(Device::gpu);
|
|
}
|
|
|
|
context.applyCommandLine(argc, argv);
|
|
return context.run();
|
|
}
|