Remove cuda conditional

framoncg · framoncg · commit 3b003ec18df4 · 2025-06-18T01:34:44.000Z
diff --git a/word_language_model/main.py b/word_language_model/main.py
@@ -264,7 +264,7 @@ def export_onnx(path, batch_size, seq_len):
     # after load the rnn params are not a continuous chunk of memory
     # this makes them a continuous chunk, and will speed up forward pass
     # Currently, only rnn model supports flatten_parameters function.
-    if args.model in ['RNN_TANH', 'RNN_RELU', 'LSTM', 'GRU'] and device.type == 'cuda':
+    if args.model in ['RNN_TANH', 'RNN_RELU', 'LSTM', 'GRU']:
         model.rnn.flatten_parameters()
 
 # Run on test data.
