Fix a small bug

Fix a small bug