allow configure gpu

src/main.rs

@@ -70,6 +70,11 @@ impl<B: Backend> Net<B> {
 }
 
 fn training() {
+	let gpu_id: usize = std::env::args()
+		.skip(1)
+		.next()
+		.map(|id| id.parse().unwrap())
+		.unwrap_or_default();
 	// load map
 	// load replay
 	// setup player
@@ -292,7 +297,7 @@ fn training() {
 		texture_data.clear();
 	}
 
-	let device = Default::default();
+	let device = burn::backend::cuda::CudaDevice::new(gpu_id);
 
 	let mut model: Net<TrainingBackend> = Net::init(&device);
 	println!("Training model ({} parameters)...", model.num_params());