From 4495e10e5159ee54814e73ae7bfb27db2754ae69 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Tue, 12 Jan 2021 20:23:40 +0800
Subject: Add Google TPU support

---
 config.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

(limited to 'config.py')

diff --git a/config.py b/config.py
index cd36cc5..19cd6f8 100644
--- a/config.py
+++ b/config.py
@@ -1,9 +1,7 @@
 config = {
     'system': {
-        # Disable accelerator
-        'disable_acc': False,
-        # GPU(s) used in training or testing if available
-        'CUDA_VISIBLE_DEVICES': '0',
+        # The number of processes for the replication
+        'nprocs': 8,
         # Directory used in training or testing for temporary storage
         'save_dir': 'runs',
     },
@@ -28,10 +26,10 @@ config = {
     },
     # Dataloader settings
     'dataloader': {
-        # Batch size (pr, k)
+        # Batch size (pr, k) * `nprocs`
         # `pr` denotes number of persons
         # `k` denotes number of sequences per person
-        'batch_size': (8, 16),
+        'batch_size': (3, 6),
         # Number of workers of Dataloader
         'num_workers': 4,
         # Faster data transfer from RAM to GPU if enabled
-- 
cgit v1.2.3