AWS STS client connection timeout while using AWS JAVA SDK

My use case requires to test connection to different AWS resources like s3 bucket, SNS topic arn and sqs url, while I am assuming a role through STS client. Now, my java application runs fine most of the times but it sometimes runs into STS client time out error and ends the API call. And when i hit the API call again, it will give the output. I wanted to know if there is any way(or any sts property) to set a custom timeout for this purpose. I am using my sts client to assume a role to a 3rd party AWS account and during the testing of the role ARN(which is to be assumed, am I getting this time out error). Please let me know if there is any way around this. Thanks.